Clean up AcademyFixedUpdateStepper when playmode changed

com.unity.ml-agents/Runtime/Academy.cs

@@ -30,6 +30,15 @@ namespace Unity.MLAgents
     /// </summary>
     internal class AcademyFixedUpdateStepper : MonoBehaviour
     {
+        void OnEnable()
+        {
+            // Check if this stepper belongs to current Academy singleton
+            // if (!Academy.Instance.IsStepperOwner(this))
+            // {
+            //     Destroy(this);
+            // }
+        }
+
         void FixedUpdate()
         {
             Academy.Instance.EnvironmentStep();
@@ -222,7 +231,25 @@ public int InferenceSeed
             Application.quitting += Dispose;
 
             LazyInitialize();
+
+#if UNITY_EDITOR
+            EditorApplication.playModeStateChanged += HandleOnPlayModeChanged;
+#endif
+        }
+
+#if UNITY_EDITOR
+        /// <summary>
+        /// Clean up the Academy when switching from edit mode to play mode
+        /// </summary>
+        /// <param name="state">State.</param>
+        void HandleOnPlayModeChanged(PlayModeStateChange state)
+        {
+            if (state == PlayModeStateChange.ExitingEditMode)
+            {
+                Dispose();
+            }
         }
+#endif
 
         /// <summary>
         /// Initialize the Academy if it hasn't already been initialized.
@@ -631,5 +658,13 @@ public void Dispose()
             // Reset the Lazy instance
             s_Lazy = new Lazy<Academy>(() => new Academy());
         }
+
+        /// <summary>
+        /// Check if the input AcademyFixedUpdateStepper belongs to this Academy.
+        /// </summary>
+        internal bool IsStepperOwner(AcademyFixedUpdateStepper stepper)
+        {
+            return stepper.gameObject.Equals(Academy.Instance.m_StepperObject);
+        }
     }
 }