Bring back root reference in grid sensor (#5300)

Project/Assets/ML-Agents/Examples/PushBlock/Prefabs/PushBlockCollabAreaGrid.prefab

   m_Modification:
     m_TransformParent: {fileID: 8188317207052398481}
     m_Modifications:
+    - target: {fileID: 1548337883655231979, guid: ac01d0f42c5e1463e943632a60d99967,
+        type: 3}
+      propertyPath: m_RootReference
+      value: 
+      objectReference: {fileID: 8190299122290044757}
+    - target: {fileID: 1548337883655231979, guid: ac01d0f42c5e1463e943632a60d99967,
+        type: 3}
+      propertyPath: m_AgentGameObject
+      value: 
+      objectReference: {fileID: 8190299122290044757}
     - target: {fileID: 2598450485826216109, guid: ac01d0f42c5e1463e943632a60d99967,
         type: 3}
       propertyPath: m_Model
     type: 3}
   m_PrefabInstance: {fileID: 6067781793364901444}
   m_PrefabAsset: {fileID: 0}
-  m_GameObject: {fileID: 0}
+  m_GameObject: {fileID: 8190299122290044757}
   m_Enabled: 1
   m_EditorHideFlags: 0
   m_Script: {fileID: 11500000, guid: d94a85eca2e074578943301959c555ba, type: 3}
     type: 3}
   m_PrefabInstance: {fileID: 6067781793364901444}
   m_PrefabAsset: {fileID: 0}
+--- !u!1 &8190299122290044757 stripped
+GameObject:
+  m_CorrespondingSourceObject: {fileID: 2710286047221272849, guid: ac01d0f42c5e1463e943632a60d99967,
+    type: 3}
+  m_PrefabInstance: {fileID: 6067781793364901444}
+  m_PrefabAsset: {fileID: 0}
 --- !u!1001 &6565363751102736699
 PrefabInstance:
   m_ObjectHideFlags: 0
     m_Modifications:
+    - target: {fileID: 1548337883655231979, guid: ac01d0f42c5e1463e943632a60d99967,
+        type: 3}
+      propertyPath: m_RootReference
+      value: 
+      objectReference: {fileID: 9115291448867436586}
+    - target: {fileID: 1548337883655231979, guid: ac01d0f42c5e1463e943632a60d99967,
+        type: 3}
+      propertyPath: m_AgentGameObject
+      value: 
+      objectReference: {fileID: 9115291448867436586}
     - target: {fileID: 2598450485826216109, guid: ac01d0f42c5e1463e943632a60d99967,
         type: 3}
       propertyPath: m_Model
     type: 3}
   m_PrefabInstance: {fileID: 6565363751102736699}
   m_PrefabAsset: {fileID: 0}
-  m_GameObject: {fileID: 0}
+  m_GameObject: {fileID: 9115291448867436586}
   m_Enabled: 1
   m_EditorHideFlags: 0
   m_Script: {fileID: 11500000, guid: d94a85eca2e074578943301959c555ba, type: 3}
     type: 3}
   m_PrefabInstance: {fileID: 6565363751102736699}
   m_PrefabAsset: {fileID: 0}
+--- !u!1 &9115291448867436586 stripped
+GameObject:
+  m_CorrespondingSourceObject: {fileID: 2710286047221272849, guid: ac01d0f42c5e1463e943632a60d99967,
+    type: 3}
+  m_PrefabInstance: {fileID: 6565363751102736699}
+  m_PrefabAsset: {fileID: 0}
 --- !u!1001 &6716844123244810954
 PrefabInstance:
   m_ObjectHideFlags: 0
     m_Modifications:
+    - target: {fileID: 1548337883655231979, guid: ac01d0f42c5e1463e943632a60d99967,
+        type: 3}
+      propertyPath: m_RootReference
+      value: 
+      objectReference: {fileID: 8695281997955662811}
+    - target: {fileID: 1548337883655231979, guid: ac01d0f42c5e1463e943632a60d99967,
+        type: 3}
+      propertyPath: m_AgentGameObject
+      value: 
+      objectReference: {fileID: 8695281997955662811}
     - target: {fileID: 2598450485826216109, guid: ac01d0f42c5e1463e943632a60d99967,
         type: 3}
       propertyPath: m_Model
     type: 3}
   m_PrefabInstance: {fileID: 6716844123244810954}
   m_PrefabAsset: {fileID: 0}
-  m_GameObject: {fileID: 0}
+  m_GameObject: {fileID: 8695281997955662811}
   m_Enabled: 1
   m_EditorHideFlags: 0
   m_Script: {fileID: 11500000, guid: d94a85eca2e074578943301959c555ba, type: 3}
     type: 3}
   m_PrefabInstance: {fileID: 6716844123244810954}
   m_PrefabAsset: {fileID: 0}
+--- !u!1 &8695281997955662811 stripped
+GameObject:
+  m_CorrespondingSourceObject: {fileID: 2710286047221272849, guid: ac01d0f42c5e1463e943632a60d99967,
+    type: 3}
+  m_PrefabInstance: {fileID: 6716844123244810954}
+  m_PrefabAsset: {fileID: 0}

Project/Assets/ML-Agents/Examples/PushBlock/Scenes/PushBlockCollab.unity

   m_Modification:
     m_TransformParent: {fileID: 0}
     m_Modifications:
+    - target: {fileID: 4704531522807670703, guid: f5bbed44a6ea747a687fbbb738eb1730,
+        type: 3}
+      propertyPath: m_ShowGizmos
+      value: 0
+      objectReference: {fileID: 0}
     - target: {fileID: 8188317207052398481, guid: f5bbed44a6ea747a687fbbb738eb1730,
         type: 3}
       propertyPath: m_RootOrder

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

                 gridSize.vector3IntValue = new Vector3Int(newGridSize.x, 1, newGridSize.z);
             }
             EditorGUI.EndDisabledGroup();
+            EditorGUILayout.PropertyField(so.FindProperty(nameof(GridSensorComponent.m_AgentGameObject)), true);
             EditorGUILayout.PropertyField(so.FindProperty(nameof(GridSensorComponent.m_RotateWithAgent)), true);
 
             EditorGUI.BeginDisabledGroup(!EditorUtilities.CanUpdateModelProperties());

com.unity.ml-agents/Runtime/Sensors/BoxOverlapChecker.cs

         Vector3Int m_GridSize;
         bool m_RotateWithAgent;
         LayerMask m_ColliderMask;
-        GameObject m_RootReference;
+        GameObject m_CenterObject;
+        GameObject m_AgentGameObject;
         string[] m_DetectableTags;
         int m_InitialColliderBufferSize;
         int m_MaxColliderBufferSize;
             Vector3Int gridSize,
             bool rotateWithAgent,
             LayerMask colliderMask,
-            GameObject rootReference,
+            GameObject centerObject,
+            GameObject agentGameObject,
             string[] detectableTags,
             int initialColliderBufferSize,
             int maxColliderBufferSize)
             m_RotateWithAgent = rotateWithAgent;
             m_ColliderMask = colliderMask;
-            m_RootReference = rootReference;
+            m_CenterObject = centerObject;
+            m_AgentGameObject = agentGameObject;
             m_DetectableTags = detectableTags;
             m_InitialColliderBufferSize = initialColliderBufferSize;
             m_MaxColliderBufferSize = maxColliderBufferSize;
         {
             if (m_RotateWithAgent)
             {
-                return m_RootReference.transform.TransformPoint(m_CellLocalPositions[cellIndex]);
+                return m_CenterObject.transform.TransformPoint(m_CellLocalPositions[cellIndex]);
-                return m_CellLocalPositions[cellIndex] + m_RootReference.transform.position;
+                return m_CellLocalPositions[cellIndex] + m_CenterObject.transform.position;
-            return m_RotateWithAgent ? m_RootReference.transform.rotation : Quaternion.identity;
+            return m_RotateWithAgent ? m_CenterObject.transform.rotation : Quaternion.identity;
         }
 
         /// <summary>
                 var currentColliderGo = foundColliders[i].gameObject;
 
                 // Continue if the current collider go is the root reference
-                if (ReferenceEquals(currentColliderGo, m_RootReference))
+                if (ReferenceEquals(currentColliderGo, m_AgentGameObject))
-                var currentDistanceSquared = (closestColliderPoint - m_RootReference.transform.position).sqrMagnitude;
+                var currentDistanceSquared = (closestColliderPoint - m_CenterObject.transform.position).sqrMagnitude;
 
                 if (currentDistanceSquared >= minDistanceSquared)
                 {
             for (int i = 0; i < numFound; i++)
             {
                 var currentColliderGo = foundColliders[i].gameObject;
-                if (!ReferenceEquals(currentColliderGo, m_RootReference))
+                if (!ReferenceEquals(currentColliderGo, m_AgentGameObject))
                 {
                     detectedAction.Invoke(currentColliderGo, cellIndex);
                 }

com.unity.ml-agents/Runtime/Sensors/GridSensorComponent.cs

         }
 
         [HideInInspector, SerializeField]
+        internal GameObject m_AgentGameObject;
+        /// <summary>
+        /// The reference of the root of the agent. This is used to disambiguate objects with
+        /// the same tag as the agent. Defaults to current GameObject.
+        /// </summary>
+        public GameObject AgentGameObject
+        {
+            get { return (m_AgentGameObject == null ? gameObject : m_AgentGameObject); }
+            set { m_AgentGameObject = value; }
+        }
+
+        [HideInInspector, SerializeField]
         internal string[] m_DetectableTags;
         /// <summary>
         /// List of tags that are detected.
                 m_RotateWithAgent,
                 m_ColliderMask,
                 gameObject,
+                AgentGameObject,
                 m_DetectableTags,
                 m_InitialColliderBufferSize,
                 m_MaxColliderBufferSize

com.unity.ml-agents/Tests/Runtime/Sensor/BoxOverlapCheckerTests.cs

             Vector3Int gridSize,
             bool rotateWithAgent,
             LayerMask colliderMask,
-            GameObject rootReference,
+            GameObject centerObject,
+            GameObject agentGameObject,
             string[] detectableTags,
             int initialColliderBufferSize,
             int maxColliderBufferSize
                 rotateWithAgent,
                 colliderMask,
-                rootReference,
+                centerObject,
+                agentGameObject,
                 detectableTags,
                 initialColliderBufferSize,
                 maxColliderBufferSize)
             int gridSizeX = 10,
             int gridSizeZ = 10,
             bool rotateWithAgent = true,
-            GameObject rootReference = null,
+            GameObject centerObject = null,
+            GameObject agentGameObject = null,
             string[] detectableTags = null,
             int initialColliderBufferSize = 4,
             int maxColliderBufferSize = 500)
                 new Vector3Int(gridSizeX, 1, gridSizeZ),
                 rotateWithAgent,
                 LayerMask.GetMask("Default"),
-                rootReference,
+                centerObject,
+                agentGameObject,
                 detectableTags,
                 initialColliderBufferSize,
                 maxColliderBufferSize);
         {
             var testGo = new GameObject("test");
             testGo.transform.position = Vector3.zero;
-            var boxOverlapSquare = TestBoxOverlapChecker.CreateChecker(gridSizeX: 10, gridSizeZ: 10, rotateWithAgent: false, rootReference: testGo);
+            var boxOverlapSquare = TestBoxOverlapChecker.CreateChecker(gridSizeX: 10, gridSizeZ: 10, rotateWithAgent: false, agentGameObject: testGo);
 
             var localPos = boxOverlapSquare.CellLocalPositions;
             Assert.AreEqual(new Vector3(-4.5f, 0, -4.5f), localPos[0]);
 
             var testGo2 = new GameObject("test");
             testGo2.transform.position = new Vector3(3.5f, 8f, 17f); // random, should have no effect on local positions
-            var boxOverlapRect = TestBoxOverlapChecker.CreateChecker(gridSizeX: 5, gridSizeZ: 15, rotateWithAgent: true, rootReference: testGo);
+            var boxOverlapRect = TestBoxOverlapChecker.CreateChecker(gridSizeX: 5, gridSizeZ: 15, rotateWithAgent: true, agentGameObject: testGo);
 
             localPos = boxOverlapRect.CellLocalPositions;
             Assert.AreEqual(new Vector3(-2f, 0, -7f), localPos[0]);
             var testGo = new GameObject("test");
             var position = new Vector3(3.5f, 8f, 17f);
             testGo.transform.position = position;
-            var boxOverlap = TestBoxOverlapChecker.CreateChecker(gridSizeX: 10, gridSizeZ: 10, rotateWithAgent: false, rootReference: testGo);
+            var boxOverlap = TestBoxOverlapChecker.CreateChecker(gridSizeX: 10, gridSizeZ: 10, rotateWithAgent: false, agentGameObject: testGo, centerObject: testGo);
 
             Assert.AreEqual(new Vector3(-4.5f, 0, -4.5f) + position, boxOverlap.GetCellGlobalPosition(0));
             Assert.AreEqual(new Vector3(-4.5f, 0, 4.5f) + position, boxOverlap.GetCellGlobalPosition(9));
             var testGo = new GameObject("test");
             var position = new Vector3(15f, 6f, 13f);
             testGo.transform.position = position;
-            var boxOverlap = TestBoxOverlapChecker.CreateChecker(gridSizeX: 5, gridSizeZ: 15, rotateWithAgent: true, rootReference: testGo);
+            var boxOverlap = TestBoxOverlapChecker.CreateChecker(gridSizeX: 5, gridSizeZ: 15, rotateWithAgent: true, agentGameObject: testGo, centerObject: testGo);
 
             Assert.AreEqual(new Vector3(-2f, 0, -7f) + position, boxOverlap.GetCellGlobalPosition(0));
             Assert.AreEqual(new Vector3(-2f, 0, 7f) + position, boxOverlap.GetCellGlobalPosition(14));
             var testGo = new GameObject("test");
             testGo.transform.position = Vector3.zero;
             testObjects.Add(testGo);
-            var boxOverlap = TestBoxOverlapChecker.CreateChecker(rootReference: testGo, initialColliderBufferSize: 2, maxColliderBufferSize: 5);
+            var boxOverlap = TestBoxOverlapChecker.CreateChecker(agentGameObject: testGo, centerObject: testGo, initialColliderBufferSize: 2, maxColliderBufferSize: 5);
             boxOverlap.Update();
             Assert.AreEqual(2, boxOverlap.ColliderBuffer.Length);
 
                 cellScaleZ: 10f,
                 gridSizeX: 2,
                 gridSizeZ: 2,
-                rootReference: testGo,
+                agentGameObject: testGo,
+                centerObject: testGo,
                 detectableTags: new [] { tag1 });
             var helper = new VerifyParseCollidersHelper();
             boxOverlap.GridOverlapDetectedClosest += helper.DetectedAction;
                 cellScaleZ: 10f,
                 gridSizeX: 2,
                 gridSizeZ: 2,
-                rootReference: testGo,
+                agentGameObject: testGo,
+                centerObject: testGo,
                 detectableTags: new [] { tag1 });
             var helper = new VerifyParseCollidersHelper();
             boxOverlap.GridOverlapDetectedAll += helper.DetectedAction;

docs/Migrating.md

 You'll need to remove the old sensor and create a new GridSensor.
 * These parameters names have changed but still refer to the same concept in the sensor: `GridNumSide` -> `GridSize`,
 `RotateToAgent` -> `RotateWithAgent`, `ObserveMask` -> `ColliderMask`, `DetectableObjects` -> `DetectableTags`
-* `RootReference` is removed and the sensor component's GameObject will always be ignored for hit results.
 * `DepthType` (`ChanelBase`/`ChannelHot`) option and `ChannelDepth` are removed. Now the default is
 one-hot encoding for detected tag. If you were using original GridSensor without overriding any method,
 switching to new GridSensor will produce similar effect for training although the actual observations