[MLA-1135] Physics sensors - optional reference body

Project/Assets/ML-Agents/Examples/Crawler/Prefabs/Crawler.prefab

@@ -2742,6 +2742,7 @@ GameObject:
   - component: {fileID: 4845971001715176662}
   - component: {fileID: 4845971001715176663}
   - component: {fileID: 4845971001715176660}
+  - component: {fileID: 4622120667686875944}
   m_Layer: 0
   m_Name: Crawler
   m_TagString: Untagged
@@ -2779,7 +2780,7 @@ MonoBehaviour:
   m_Name: 
   m_EditorClassIdentifier: 
   m_BrainParameters:
-    VectorObservationSize: 138
+    VectorObservationSize: 21
     NumStackedVectorObservations: 1
     VectorActionSize: 14000000
     VectorActionDescriptions: []
@@ -2872,6 +2873,30 @@ MonoBehaviour:
   m_Name: 
   m_EditorClassIdentifier: 
   debugCommandLineOverride: 
+--- !u!114 &4622120667686875944
+MonoBehaviour:
+  m_ObjectHideFlags: 0
+  m_CorrespondingSourceObject: {fileID: 0}
+  m_PrefabInstance: {fileID: 0}
+  m_PrefabAsset: {fileID: 0}
+  m_GameObject: {fileID: 4845971001715176661}
+  m_Enabled: 1
+  m_EditorHideFlags: 0
+  m_Script: {fileID: 11500000, guid: df0f8be9a37d6486498061e2cbc4cd94, type: 3}
+  m_Name: 
+  m_EditorClassIdentifier: 
+  RootBody: {fileID: 4845971001588102145}
+  VirtualRoot: {fileID: 2270141184585723037}
+  Settings:
+    UseModelSpaceTranslations: 1
+    UseModelSpaceRotations: 1
+    UseLocalSpaceTranslations: 0
+    UseLocalSpaceRotations: 1
+    UseModelSpaceLinearVelocity: 1
+    UseLocalSpaceLinearVelocity: 0
+    UseJointPositionsAndAngles: 0
+    UseJointForces: 0
+  sensorName: 
 --- !u!1 &4845971001730692034
 GameObject:
   m_ObjectHideFlags: 0
@@ -3018,6 +3043,12 @@ PrefabInstance:
       objectReference: {fileID: 0}
     m_RemovedComponents: []
   m_SourcePrefab: {fileID: 100100000, guid: 72f745913c5a34df5aaadd5c1f0024cb, type: 3}
+--- !u!1 &2270141184585723037 stripped
+GameObject:
+  m_CorrespondingSourceObject: {fileID: 2591864627249999519, guid: 72f745913c5a34df5aaadd5c1f0024cb,
+    type: 3}
+  m_PrefabInstance: {fileID: 4357529801223143938}
+  m_PrefabAsset: {fileID: 0}
 --- !u!4 &2270141184585723026 stripped
 Transform:
   m_CorrespondingSourceObject: {fileID: 2591864627249999504, guid: 72f745913c5a34df5aaadd5c1f0024cb,
@@ -3030,7 +3061,7 @@ MonoBehaviour:
     type: 3}
   m_PrefabInstance: {fileID: 4357529801223143938}
   m_PrefabAsset: {fileID: 0}
-  m_GameObject: {fileID: 0}
+  m_GameObject: {fileID: 2270141184585723037}
   m_Enabled: 1
   m_EditorHideFlags: 0
   m_Script: {fileID: 11500000, guid: 771e78c5e980e440e8cd19716b55075f, type: 3}

Project/Assets/ML-Agents/Examples/Crawler/Prefabs/FixedPlatform.prefab

@@ -392,6 +392,11 @@ PrefabInstance:
       propertyPath: targetToLookAt
       value: 
       objectReference: {fileID: 2673081981996998229}
+    - target: {fileID: 4622120667686875944, guid: 0456c89e8c9c243d595b039fe7aa0bf9,
+        type: 3}
+      propertyPath: Settings.UseLocalSpaceLinearVelocity
+      value: 1
+      objectReference: {fileID: 0}
     - target: {fileID: 4845971000000621469, guid: 0456c89e8c9c243d595b039fe7aa0bf9,
         type: 3}
       propertyPath: m_ConnectedAnchor.x

Project/Assets/ML-Agents/Examples/Crawler/Scripts/CrawlerAgent.cs

@@ -91,17 +91,8 @@ public void CollectObservationBodyPart(BodyPart bp, VectorSensor sensor)
         //GROUND CHECK
         sensor.AddObservation(bp.groundContact.touchingGround); // Is this bp touching the ground
 
-        //Get velocities in the context of our orientation cube's space
-        //Note: You can get these velocities in world space as well but it may not train as well.
-        sensor.AddObservation(orientationCube.transform.InverseTransformDirection(bp.rb.velocity));
-        sensor.AddObservation(orientationCube.transform.InverseTransformDirection(bp.rb.angularVelocity));
-
-        //Get position relative to hips in the context of our orientation cube's space
-        sensor.AddObservation(orientationCube.transform.InverseTransformDirection(bp.rb.position - body.position));
-
         if (bp.rb.transform != body)
         {
-            sensor.AddObservation(bp.rb.transform.localRotation);
             sensor.AddObservation(bp.currentStrength / m_JdController.maxJointForceLimit);
         }
     }
@@ -111,9 +102,6 @@ public void CollectObservationBodyPart(BodyPart bp, VectorSensor sensor)
     /// </summary>
     public override void CollectObservations(VectorSensor sensor)
     {
-        //Add body rotation delta relative to orientation cube
-        sensor.AddObservation(Quaternion.FromToRotation(body.forward, orientationCube.transform.forward));
-
         //Add pos of target relative to orientation cube
         sensor.AddObservation(orientationCube.transform.InverseTransformPoint(target.transform.position));
 

com.unity.ml-agents.extensions/README.md

@@ -1,3 +1,5 @@
 # ML-Agents Extensions
 
 This is a source-only package for new features based on ML-Agents.
+
+More details coming soon.

com.unity.ml-agents.extensions/Runtime/AssemblyInfo.cs

@@ -0,0 +1,3 @@
+using System.Runtime.CompilerServices;
+
+[assembly: InternalsVisibleTo("Unity.ML-Agents.Extensions.EditorTests")]

com.unity.ml-agents.extensions/Runtime/Sensors/ArticulationBodyPoseExtractor.cs

@@ -54,17 +54,17 @@ public ArticulationBodyPoseExtractor(ArticulationBody rootBody)
                 parentIndices[i] = bodyToIndex[parentArticBody];
             }
 
-            SetParentIndices(parentIndices);
+            Setup(parentIndices);
         }
 
         /// <inheritdoc/>
-        protected override Vector3 GetLinearVelocityAt(int index)
+        protected internal override Vector3 GetLinearVelocityAt(int index)
         {
             return m_Bodies[index].velocity;
         }
 
         /// <inheritdoc/>
-        protected override Pose GetPoseAt(int index)
+        protected internal override Pose GetPoseAt(int index)
         {
             var body = m_Bodies[index];
             var go = body.gameObject;

com.unity.ml-agents.extensions/Runtime/Sensors/PhysicsBodySensor.cs

@@ -18,12 +18,20 @@ public class PhysicsBodySensor : ISensor
         /// <summary>
         ///  Construct a new PhysicsBodySensor
         /// </summary>
-        /// <param name="rootBody"></param>
+        /// <param name="rootBody">The root Rigidbody. This has no Joints on it (but other Joints may connect to it).</param>
+        /// <param name="rootGameObject">Optional GameObject used to find Rigidbodies in the hierarchy.</param>
+        /// <param name="virtualRoot">Optional GameObject used to determine the root of the poses,
         /// <param name="settings"></param>
         /// <param name="sensorName"></param>
-        public PhysicsBodySensor(Rigidbody rootBody, GameObject rootGameObject, PhysicsSensorSettings settings, string sensorName=null)
+        public PhysicsBodySensor(
+            Rigidbody rootBody,
+            GameObject rootGameObject,
+            GameObject virtualRoot,
+            PhysicsSensorSettings settings,
+            string sensorName=null
+        )
         {
-            var poseExtractor = new RigidBodyPoseExtractor(rootBody, rootGameObject);
+            var poseExtractor = new RigidBodyPoseExtractor(rootBody, rootGameObject, virtualRoot);
             m_PoseExtractor = poseExtractor;
             m_SensorName = string.IsNullOrEmpty(sensorName) ? $"PhysicsBodySensor:{rootBody?.name}" : sensorName;
             m_Settings = settings;

com.unity.ml-agents.extensions/Runtime/Sensors/PhysicsSensorSettings.cs

@@ -1,5 +1,4 @@
 using System;
-
 using Unity.MLAgents.Sensors;
 
 namespace Unity.MLAgents.Extensions.Sensors
@@ -95,51 +94,53 @@ public static int WritePoses(this ObservationWriter writer, PhysicsSensorSetting
             var offset = baseOffset;
             if (settings.UseModelSpace)
             {
-                var poses = poseExtractor.ModelSpacePoses;
-                var vels = poseExtractor.ModelSpaceVelocities;
-
-                for(var i=0; i<poseExtractor.NumPoses; i++)
+                foreach (var pose in poseExtractor.GetEnabledModelSpacePoses())
                 {
-                    var pose = poses[i];
-                    if(settings.UseModelSpaceTranslations)
+                    if (settings.UseModelSpaceTranslations)
                     {
                         writer.Add(pose.position, offset);
                         offset += 3;
                     }
+
                     if (settings.UseModelSpaceRotations)
                     {
                         writer.Add(pose.rotation, offset);
                         offset += 4;
                     }
+                }
+
+                foreach(var vel in poseExtractor.GetEnabledModelSpaceVelocities())
+                {
                     if (settings.UseModelSpaceLinearVelocity)
                     {
-                        writer.Add(vels[i], offset);
+                        writer.Add(vel, offset);
                         offset += 3;
                     }
                 }
             }
 
             if (settings.UseLocalSpace)
             {
-                var poses = poseExtractor.LocalSpacePoses;
-                var vels = poseExtractor.LocalSpaceVelocities;
-
-                for(var i=0; i<poseExtractor.NumPoses; i++)
+                foreach (var pose in poseExtractor.GetEnabledLocalSpacePoses())
                 {
-                    var pose = poses[i];
-                    if(settings.UseLocalSpaceTranslations)
+                    if (settings.UseLocalSpaceTranslations)
                     {
                         writer.Add(pose.position, offset);
                         offset += 3;
                     }
+
                     if (settings.UseLocalSpaceRotations)
                     {
                         writer.Add(pose.rotation, offset);
                         offset += 4;
                     }
+                }
+
+                foreach(var vel in poseExtractor.GetEnabledLocalSpaceVelocities())
+                {
                     if (settings.UseLocalSpaceLinearVelocity)
                     {
-                        writer.Add(vels[i], offset);
+                        writer.Add(vel, offset);
                         offset += 3;
                     }
                 }