[MLA-1474] detect recursion on Agent methods and throw

com.unity.ml-agents/CHANGELOG.md

@@ -18,6 +18,9 @@ and this project adheres to
 
 ### Bug Fixes
 #### com.unity.ml-agents (C#)
+- `Agent.CollectObservations()` and `Agent.EndEpisode()` will now throw an exception
+if they are called recursively (for example, if they call `Agent.EndEpisode()`).
+Previously, this would result in an infinite loop and cause the editor to hang. (#4573)
 #### ml-agents / ml-agents-envs / gym-unity (Python)
 
 

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

@@ -155,9 +155,8 @@ public bool IsCommunicatorOn
         // Flag used to keep track of the first time the Academy is reset.
         bool m_HadFirstReset;
 
-        // Whether the Academy is in the middle of a step. This is used to detect and Academy
-        // step called by user code that is also called by the Academy.
-        bool m_IsStepping;
+        // Detect an Academy step called by user code that is also called by the Academy.
+        private RecursionChecker m_StepRecursionChecker = new RecursionChecker("EnvironmentStep");
 
         // Random seed used for inference.
         int m_InferenceSeed;
@@ -535,22 +534,7 @@ void ForcedFullReset()
         /// </summary>
         public void EnvironmentStep()
         {
-            // Check whether we're already in the middle of a step.
-            // This shouldn't happen generally, but could happen if user code (e.g. CollectObservations)
-            // that is called by EnvironmentStep() also calls EnvironmentStep(). This would result
-            // in an infinite loop and/or stack overflow, so stop it before it happens.
-            if (m_IsStepping)
-            {
-                throw new UnityAgentsException(
-                    "Academy.EnvironmentStep() called recursively. " +
-                    "This might happen if you call EnvironmentStep() from custom code such as " +
-                    "CollectObservations() or OnActionReceived()."
-                );
-            }
-
-            m_IsStepping = true;
-
-            try
+            using (m_StepRecursionChecker.Start())
             {
                 if (!m_HadFirstReset)
                 {
@@ -584,11 +568,6 @@ public void EnvironmentStep()
                     AgentAct?.Invoke();
                 }
             }
-            finally
-            {
-                // Reset m_IsStepping when we're done (or if an exception occurred).
-                m_IsStepping = false;
-            }
         }
 
         /// <summary>

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

@@ -283,6 +283,9 @@ internal struct AgentParameters
         /// </summary>
         internal VectorSensor collectObservationsSensor;
 
+        private RecursionChecker m_CollectObservationsChecker = new RecursionChecker("CollectObservations");
+        private RecursionChecker m_OnEpisodeBeginChecker = new RecursionChecker("OnEpisodeBegin");
+
         /// <summary>
         /// List of IActuators that this Agent will delegate actions to if any exist.
         /// </summary>
@@ -435,7 +438,10 @@ public void LazyInitialize()
             // episode when initializing until after the Academy had its first reset.
             if (Academy.Instance.TotalStepCount != 0)
             {
-                OnEpisodeBegin();
+                using (m_OnEpisodeBeginChecker.Start())
+                {
+                    OnEpisodeBegin();
+                }
             }
         }
 
@@ -512,7 +518,10 @@ void NotifyAgentDone(DoneReason doneReason)
             {
                 // Make sure the latest observations are being passed to training.
                 collectObservationsSensor.Reset();
-                CollectObservations(collectObservationsSensor);
+                using (m_CollectObservationsChecker.Start())
+                {
+                    CollectObservations(collectObservationsSensor);
+                }
             }
             // Request the last decision with no callbacks
             // We request a decision so Python knows the Agent is done immediately
@@ -1006,7 +1015,10 @@ void SendInfoToBrain()
             UpdateSensors();
             using (TimerStack.Instance.Scoped("CollectObservations"))
             {
-                CollectObservations(collectObservationsSensor);
+                using (m_CollectObservationsChecker.Start())
+                {
+                    CollectObservations(collectObservationsSensor);
+                }
             }
             using (TimerStack.Instance.Scoped("CollectDiscreteActionMasks"))
             {
@@ -1229,7 +1241,11 @@ void _AgentReset()
         {
             ResetData();
             m_StepCount = 0;
-            OnEpisodeBegin();
+            using (m_OnEpisodeBeginChecker.Start())
+            {
+                OnEpisodeBegin();
+            }
+
         }
 
         /// <summary>

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

@@ -0,0 +1,35 @@
+using System;
+
+namespace Unity.MLAgents
+{
+    internal class RecursionChecker : IDisposable
+    {
+        private bool m_IsRunning;
+        private string m_MethodName;
+
+        public RecursionChecker(string methodName)
+        {
+            m_MethodName = methodName;
+        }
+
+        public IDisposable Start()
+        {
+            if (m_IsRunning)
+            {
+                throw new UnityAgentsException(
+                    $"{m_MethodName} called recursively. " +
+                    "This might happen if you call EnvironmentStep() or EndEpisode() from custom " +
+                    "code such as CollectObservations() or OnActionReceived()."
+                );
+            }
+            m_IsRunning = true;
+            return this;
+        }
+
+        public void Dispose()
+        {
+            // Reset the flag when we're done (or if an exception occurred).
+            m_IsRunning = false;
+        }
+    }
+}

com.unity.ml-agents/Tests/Editor/MLAgentsEditModeTest.cs

@@ -830,4 +830,60 @@ public void TestObservableAttributeBehaviorIgnore()
             }
         }
     }
+
+    [TestFixture]
+    public class AgentRecursionTests
+    {
+        [SetUp]
+        public void SetUp()
+        {
+            if (Academy.IsInitialized)
+            {
+                Academy.Instance.Dispose();
+            }
+        }
+
+        class CollectObsEndEpisodeAgent : Agent
+        {
+            public override void CollectObservations(VectorSensor sensor)
+            {
+                // NEVER DO THIS IN REAL CODE!
+                EndEpisode();
+            }
+        }
+
+        class OnEpisodeBeginEndEpisodeAgent : Agent
+        {
+            public override void OnEpisodeBegin()
+            {
+                // NEVER DO THIS IN REAL CODE!
+                EndEpisode();
+            }
+        }
+
+        void TestRecursiveThrows<T>() where T : Agent
+        {
+            var gameObj = new GameObject();
+            var agent = gameObj.AddComponent<T>();
+            agent.LazyInitialize();
+            agent.RequestDecision();
+
+            Assert.Throws<UnityAgentsException>(() =>
+            {
+                Academy.Instance.EnvironmentStep();
+            });
+        }
+
+        [Test]
+        public void TestRecursiveCollectObsEndEpisodeThrows()
+        {
+            TestRecursiveThrows<CollectObsEndEpisodeAgent>();
+        }
+
+        [Test]
+        public void TestRecursiveOnEpisodeBeginEndEpisodeThrows()
+        {
+            TestRecursiveThrows<OnEpisodeBeginEndEpisodeAgent>();
+        }
+    }
 }

com.unity.ml-agents/Tests/Editor/RecursionCheckerTests.cs

@@ -0,0 +1,72 @@
+using System;
+using NUnit.Framework;
+
+namespace Unity.MLAgents.Tests
+{
+    [TestFixture]
+    public class RecursionCheckerTests
+    {
+        class InfiniteRecurser
+        {
+            RecursionChecker m_checker = new RecursionChecker("InfiniteRecurser");
+            public int NumCalls = 0;
+
+            public void Implode()
+            {
+                NumCalls++;
+                using (m_checker.Start())
+                {
+                    Implode();
+                }
+            }
+        }
+
+        [Test]
+        public void TestRecursionCheck()
+        {
+            var rc = new InfiniteRecurser();
+            Assert.Throws<UnityAgentsException>(() =>
+            {
+                rc.Implode();
+            });
+
+            // Should increment twice before bailing out.
+            Assert.AreEqual(2, rc.NumCalls);
+        }
+
+        class OneTimeThrower
+        {
+            RecursionChecker m_checker = new RecursionChecker("OneTimeThrower");
+            public int NumCalls;
+
+            public void DoStuff()
+            {
+                // This method throws from inside the checker the first time.
+                // Later calls do nothing.
+                NumCalls++;
+                using (m_checker.Start())
+                {
+                    if (NumCalls == 1)
+                    {
+                        throw new ArgumentException("oops");
+                    }
+                }
+            }
+        }
+
+        [Test]
+        public void TestThrowResetsFlag()
+        {
+            var ott = new OneTimeThrower();
+            Assert.Throws<ArgumentException>(() =>
+            {
+                ott.DoStuff();
+            });
+
+            // Make sure the flag is cleared if we throw in the "using". Should be able to step subsequently.
+            ott.DoStuff();
+            ott.DoStuff();
+            Assert.AreEqual(3, ott.NumCalls);
+        }
+    }
+}