diff --git a/com.unity.ml-agents/Runtime/Academy.cs b/com.unity.ml-agents/Runtime/Academy.cs
index 8ed370e350..991bdaee53 100644
--- a/com.unity.ml-agents/Runtime/Academy.cs
+++ b/com.unity.ml-agents/Runtime/Academy.cs
@@ -25,6 +25,26 @@
namespace Unity.MLAgents
{
+ internal class RecursionChecker : IDisposable
+ {
+ private bool m_IsRunning;
+
+ public IDisposable Start()
+ {
+ if (m_IsRunning)
+ {
+ throw new UnityAgentsException("Don't do this.");
+ }
+ m_IsRunning = true;
+ return this;
+ }
+
+ public void Dispose()
+ {
+ m_IsRunning = false;
+ }
+ }
+
///
/// Helper class to step the Academy during FixedUpdate phase.
///
@@ -157,7 +177,7 @@ public bool IsCommunicatorOn
// 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;
+ private RecursionChecker m_StepRecursionChecker = new RecursionChecker();
// Random seed used for inference.
int m_InferenceSeed;
@@ -535,22 +555,7 @@ void ForcedFullReset()
///
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 +589,6 @@ public void EnvironmentStep()
AgentAct?.Invoke();
}
}
- finally
- {
- // Reset m_IsStepping when we're done (or if an exception occurred).
- m_IsStepping = false;
- }
}
///