Remove RB sensor and add X Z velocity obs

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_RootReference)), 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_CenterObject;
         GameObject m_RootReference;
         string[] m_DetectableTags;
         int m_InitialColliderBufferSize;
             Vector3Int gridSize,
             bool rotateWithAgent,
             LayerMask colliderMask,
+            GameObject centerObject,
             GameObject rootReference,
             string[] detectableTags,
             int initialColliderBufferSize,
             m_GridSize = gridSize;
             m_RotateWithAgent = rotateWithAgent;
             m_ColliderMask = colliderMask;
+            m_CenterObject = centerObject;
             m_RootReference = rootReference;
             m_DetectableTags = detectableTags;
             m_InitialColliderBufferSize = initialColliderBufferSize;
         {
             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 closestColliderPoint = foundColliders[i].ClosestPointOnBounds(cellCenter);
-                var currentDistanceSquared = (closestColliderPoint - m_RootReference.transform.position).sqrMagnitude;
+                var currentDistanceSquared = (closestColliderPoint - m_CenterObject.transform.position).sqrMagnitude;
 
                 if (currentDistanceSquared >= minDistanceSquared)
                 {

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

         }
 
         [HideInInspector, SerializeField]
+        internal GameObject m_RootReference;
+        /// <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 RootReference
+        {
+            get { return (m_RootReference == null ? gameObject : m_RootReference); }
+            set { m_RootReference = value; }
+        }
+
+        [HideInInspector, SerializeField]
         internal string[] m_DetectableTags;
         /// <summary>
         /// List of tags that are detected.
                 m_RotateWithAgent,
                 m_ColliderMask,
                 gameObject,
+                RootReference,
                 m_DetectableTags,
                 m_InitialColliderBufferSize,
                 m_MaxColliderBufferSize

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

             Vector3Int gridSize,
             bool rotateWithAgent,
             LayerMask colliderMask,
+            GameObject centerObject,
             GameObject rootReference,
             string[] detectableTags,
             int initialColliderBufferSize,
                 gridSize,
                 rotateWithAgent,
                 colliderMask,
+                centerObject,
                 rootReference,
                 detectableTags,
                 initialColliderBufferSize,
             int gridSizeX = 10,
             int gridSizeZ = 10,
             bool rotateWithAgent = true,
+            GameObject centerObject = null,
             GameObject rootReference = null,
             string[] detectableTags = null,
             int initialColliderBufferSize = 4,
                 new Vector3Int(gridSizeX, 1, gridSizeZ),
                 rotateWithAgent,
                 LayerMask.GetMask("Default"),
+                centerObject,
                 rootReference,
                 detectableTags,
                 initialColliderBufferSize,
             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, rootReference: 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, rootReference: 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(rootReference: testGo, centerObject: testGo, initialColliderBufferSize: 2, maxColliderBufferSize: 5);
             boxOverlap.Update();
             Assert.AreEqual(2, boxOverlap.ColliderBuffer.Length);
 
                 gridSizeX: 2,
                 gridSizeZ: 2,
                 rootReference: testGo,
+                centerObject: testGo,
                 detectableTags: new [] { tag1 });
             var helper = new VerifyParseCollidersHelper();
             boxOverlap.GridOverlapDetectedClosest += helper.DetectedAction;
                 gridSizeX: 2,
                 gridSizeZ: 2,
                 rootReference: testGo,
+                centerObject: testGo,
                 detectableTags: new [] { tag1 });
             var helper = new VerifyParseCollidersHelper();
             boxOverlap.GridOverlapDetectedAll += helper.DetectedAction;

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: 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: 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: 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/Prefabs/PushBlockAgentGridCollab.prefab

   m_CellScale: {x: 1, y: 0.5, z: 1}
   m_GridSize: {x: 20, y: 1, z: 20}
   m_RotateWithAgent: 1
+  m_RootReference: {fileID: 0}
   m_DetectableTags:
   - wall
   - agent
   - component: {fileID: 2698669610236060893}
   - component: {fileID: 6481017652691923713}
   - component: {fileID: 2106090979146236401}
-  - component: {fileID: 294674626881224982}
   m_Layer: 0
   m_Name: PushBlockAgentGridCollab
   m_TagString: agent
   m_Name: 
   m_EditorClassIdentifier: 
   m_BrainParameters:
-    VectorObservationSize: 0
-    NumStackedVectorObservations: 2
+    VectorObservationSize: 2
+    NumStackedVectorObservations: 1
     m_ActionSpec:
       m_NumContinuousActions: 0
       BranchSizes: 07000000
   m_Name: 
   m_EditorClassIdentifier: 
   debugCommandLineOverride: 
---- !u!114 &294674626881224982
-MonoBehaviour:
-  m_ObjectHideFlags: 0
-  m_CorrespondingSourceObject: {fileID: 0}
-  m_PrefabInstance: {fileID: 0}
-  m_PrefabAsset: {fileID: 0}
-  m_GameObject: {fileID: 2710286047221272849}
-  m_Enabled: 1
-  m_EditorHideFlags: 0
-  m_Script: {fileID: 11500000, guid: df0f8be9a37d6486498061e2cbc4cd94, type: 3}
-  m_Name: 
-  m_EditorClassIdentifier: 
-  RootBody: {fileID: 2691837803371243691}
-  VirtualRoot: {fileID: 2710286047221272849}
-  Settings:
-    UseModelSpaceTranslations: 0
-    UseModelSpaceRotations: 0
-    UseLocalSpaceTranslations: 1
-    UseLocalSpaceRotations: 1
-    UseModelSpaceLinearVelocity: 1
-    UseLocalSpaceLinearVelocity: 0
-    UseJointPositionsAndAngles: 0
-    UseJointForces: 0
-  sensorName: RBSensor
 --- !u!1 &2710572731397926669
 GameObject:
   m_ObjectHideFlags: 0

Project/Assets/ML-Agents/Examples/PushBlock/Scripts/PushAgentCollab.cs

 using UnityEngine;
 using Unity.MLAgents;
 using Unity.MLAgents.Actuators;
+using Unity.MLAgents.Sensors;
 
 public class PushAgentCollab : Agent
 {
     {
         // Move the agent using the action.
         MoveAgent(actionBuffers.DiscreteActions);
+    }
+
+    /// <summary>
+    /// Called every time a decision is requested.
+    /// </summary>
+    public override void CollectObservations(VectorSensor sensor)
+    {
+        sensor.AddObservation(Vector3.Dot(m_AgentRb.velocity, m_AgentRb.transform.forward));
+        sensor.AddObservation(Vector3.Dot(m_AgentRb.velocity, m_AgentRb.transform.right));
     }
 
     public override void Heuristic(in ActionBuffers actionsOut)