update PushBlockTeamManager

Project/Assets/ML-Agents/Examples/PushBlock/Scripts/PushBlockEnvController.cs

 using System.Collections;
 using System.Collections.Generic;
-using Unity.MLAgents;
+using Unity.MLAgents.Extensions.Teams;
 using UnityEngine;
 
 public class PushBlockEnvController : MonoBehaviour
     /// </summary>
     /// <returns></returns>
     [Header("Max Environment Steps")] public int MaxEnvironmentSteps = 25000;
-    private int m_ResetTimer;
+    // private int m_ResetTimer;
 
     /// <summary>
     /// The area bounds.
             item.Rb = item.T.GetComponent<Rigidbody>();
         }
         // Initialize TeamManager
-        m_TeamManager = new PushBlockTeamManager();
+        m_TeamManager = new PushBlockTeamManager(this);
+        m_TeamManager.SetTeamMaxStep(MaxEnvironmentSteps);
         foreach (var item in AgentsList)
         {
             item.StartingPos = item.Agent.transform.position;
 
     }
 
-    // Update is called once per frame
-    void FixedUpdate()
-    {
-        m_ResetTimer += 1;
-        if (m_ResetTimer > MaxEnvironmentSteps)
-        {
-            ResetScene();
-        }
-    }
-
     /// <summary>
     /// Use the ground's bounds to pick a random spawn position.
     /// </summary>
         if (done)
         {
             //Reset assets
-            ResetScene();
+            m_TeamManager.EndTeamEpisode();
         }
     }
 
     }
 
-    void ResetScene()
+    public void ResetScene()
-        m_ResetTimer = 0;
+        // m_ResetTimer = 0;
 
         //Random platform rot
         var rotation = Random.Range(0, 4);
         //End Episode
-        foreach (var item in AgentsList)
-        {
-            if (!item.Agent)
-            {
-                return;
-            }
-            item.Agent.EndEpisode();
-        }
+        // foreach (var item in AgentsList)
+        // {
+        //     if (!item.Agent)
+        //     {
+        //         return;
+        //     }
+        //     item.Agent.EndEpisode();
+        // }
         //Reset Agents
         foreach (var item in AgentsList)
         {

Project/Assets/ML-Agents/Examples/PushBlock/Scripts/PushBlockTeamManager.cs

 
 public class PushBlockTeamManager : BaseTeamManager
 {
-    List<Agent> m_AgentList = new List<Agent> { };
+    PushBlockEnvController m_envController;
-    public override void RegisterAgent(Agent agent)
+    public PushBlockTeamManager(PushBlockEnvController envController)
+    {
+        m_envController = envController;
+    }
+
+    public override void OnTeamEpisodeBegin()
-        m_AgentList.Add(agent);
+        m_envController.ResetScene();
     }
 }

Project/Assets/ML-Agents/Examples/PushBlock/Scripts/ZombiePushBlockDeathEnvController.cs

             item.Rb = item.T.GetComponent<Rigidbody>();
         }
         // Initialize TeamManager
-        m_TeamManager = new PushBlockTeamManager();
+        // m_TeamManager = new PushBlockTeamManager();
         foreach (var item in AgentsList)
         {
             item.StartingPos = item.Agent.transform.position;
             // Add to team manager
-            item.Agent.SetTeamManager(m_TeamManager);
+            // item.Agent.SetTeamManager(m_TeamManager);
         }
         foreach (var item in ZombiesList)
         {