Fix serialization of RigidBodySensorComponents

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;
@@ -47,13 +46,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;
@@ -69,7 +68,5 @@ public override void OnInspectorGUI()
                 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.

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

@@ -94,7 +94,7 @@ internal void SetPoseEnabled(int index, bool enabled)
 
         internal bool IsTrivial()
         {
-            if (ReferenceEquals(RootBody, null))
+            if (ReferenceEquals(RootBody, null) || !RootBody)
             {
                 // It *is* trivial, but this will happen when the sensor is being set up, so don't warn then.
                 return false;