Fix serialization of RigidBodySensorComponents

The settings of `RigidBodySensorComponent`s were not getting serialized
and saved properly before. Now they are.
Fixes Unity-Technologies#5774

com.unity.ml-agents.extensions/Editor/RigidBodySensorComponentEditor.cs

@@ -1,11 +1,10 @@
-using UnityEditor;
 using Unity.MLAgents.Editor;
 using Unity.MLAgents.Extensions.Sensors;
+using UnityEditor;
 
 namespace Unity.MLAgents.Extensions.Editor
 {
     [CustomEditor(typeof(RigidBodySensorComponent))]
-    [CanEditMultipleObjects]
     internal class RigidBodySensorComponentEditor : UnityEditor.Editor
     {
         bool ShowHierarchy = true;
@@ -26,8 +25,6 @@ public override void OnInspectorGUI()
                 );
             }
 
-            bool requireExtractorUpdate;
-
             EditorGUI.BeginDisabledGroup(!EditorUtilities.CanUpdateModelProperties());
             {
                 // All the fields affect the sensor order or observation size,
@@ -36,8 +33,11 @@ public override void OnInspectorGUI()
                 EditorGUILayout.PropertyField(so.FindProperty("RootBody"), true);
                 EditorGUILayout.PropertyField(so.FindProperty("VirtualRoot"), true);
 
-                // Changing the root body or virtual root changes the hierarchy, so we need to reset later.
-                requireExtractorUpdate = EditorGUI.EndChangeCheck();
+                // Changing the root body or virtual root changes the hierarchy, so we need to reset.
+                if (EditorGUI.EndChangeCheck())
+                {
+                    rbSensorComp.ResetPoseExtractor();
+                }
 
                 EditorGUILayout.PropertyField(so.FindProperty("Settings"), true);
 
@@ -47,13 +47,13 @@ public override void OnInspectorGUI()
                 {
                     var treeNodes = rbSensorComp.GetDisplayNodes();
                     var originalIndent = EditorGUI.indentLevel;
+                    var poseEnabled = so.FindProperty("m_PoseExtractor").FindPropertyRelative("m_PoseEnabled");
                     foreach (var node in treeNodes)
                     {
                         var obj = node.NodeObject;
                         var objContents = EditorGUIUtility.ObjectContent(obj, obj.GetType());
                         EditorGUI.indentLevel = originalIndent + node.Depth;
-                        var enabled = EditorGUILayout.Toggle(objContents, node.Enabled);
-                        rbSensorComp.SetPoseEnabled(node.OriginalIndex, enabled);
+                        EditorGUILayout.PropertyField(poseEnabled.GetArrayElementAtIndex(node.OriginalIndex), objContents);
                     }
 
                     EditorGUI.indentLevel = originalIndent;
@@ -64,12 +64,6 @@ public override void OnInspectorGUI()
             EditorGUI.EndDisabledGroup();
 
             so.ApplyModifiedProperties();
-            if (requireExtractorUpdate)
-            {
-                rbSensorComp.ResetPoseExtractor();
-            }
         }
-
-
     }
 }

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

@@ -14,16 +14,17 @@ namespace Unity.MLAgents.Extensions.Sensors
     /// Poses are either considered in model space, which is relative to a root body,
     /// or in local space, which is relative to their parent.
     /// </summary>
-    public abstract class PoseExtractor
+    [Serializable]
+    public abstract class PoseExtractor : ISerializationCallbackReceiver
     {
-        int[] m_ParentIndices;
+        [SerializeField] int[] m_ParentIndices;
         Pose[] m_ModelSpacePoses;
         Pose[] m_LocalSpacePoses;
 
         Vector3[] m_ModelSpaceLinearVelocities;
         Vector3[] m_LocalSpaceLinearVelocities;
 
-        bool[] m_PoseEnabled;
+        [SerializeField] bool[] m_PoseEnabled;
 
 
         /// <summary>
@@ -177,11 +178,8 @@ protected void Setup(int[] parentIndices)
 #endif
             m_ParentIndices = parentIndices;
             var numPoses = parentIndices.Length;
-            m_ModelSpacePoses = new Pose[numPoses];
-            m_LocalSpacePoses = new Pose[numPoses];
 
-            m_ModelSpaceLinearVelocities = new Vector3[numPoses];
-            m_LocalSpaceLinearVelocities = new Vector3[numPoses];
+            CreateSensorArrays(numPoses);
 
             m_PoseEnabled = new bool[numPoses];
             // All poses are enabled by default. Generally we'll want to disable the root though.
@@ -191,6 +189,25 @@ protected void Setup(int[] parentIndices)
             }
         }
 
+        public void OnBeforeSerialize()
+        {
+            // no-op
+        }
+
+        public void OnAfterDeserialize()
+        {
+            CreateSensorArrays(m_ParentIndices.Length);
+        }
+
+        private void CreateSensorArrays(int numPoses)
+        {
+            m_ModelSpacePoses = new Pose[numPoses];
+            m_LocalSpacePoses = new Pose[numPoses];
+
+            m_ModelSpaceLinearVelocities = new Vector3[numPoses];
+            m_LocalSpaceLinearVelocities = new Vector3[numPoses];
+        }
+
         /// <summary>
         /// Return the world space Pose of the i'th object.
         /// </summary>

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

@@ -1,5 +1,7 @@
+using System;
 using System.Collections.Generic;
 using UnityEngine;
+using Object = UnityEngine.Object;
 
 namespace Unity.MLAgents.Extensions.Sensors
 {
@@ -8,15 +10,16 @@ namespace Unity.MLAgents.Extensions.Sensors
     /// Utility class to track a hierarchy of RigidBodies. These are assumed to have a root node,
     /// and child nodes are connect to their parents via Joints.
     /// </summary>
+    [Serializable]
     public class RigidBodyPoseExtractor : PoseExtractor
     {
-        Rigidbody[] m_Bodies;
+        [SerializeField] Rigidbody[] m_Bodies;
 
         /// <summary>
         /// Optional game object used to determine the root of the poses, separate from the actual Rigidbodies
         /// in the hierarchy. For locomotion
         /// </summary>
-        GameObject m_VirtualRoot;
+        [SerializeField] GameObject m_VirtualRoot;
 
         /// <summary>
         /// Initialize given a root RigidBody.