Freeze agent, don't kill it

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

     //    Rigidbody m_BlockRb;  //cached on initialization
     Rigidbody m_AgentRb;  //cached on initialization
 
+    public bool frozen = false;
+
     public bool useVectorObs = true;
     void Awake()
     {
     public override void OnActionReceived(ActionBuffers actionBuffers)
 
     {
-        // Move the agent using the action.
-        MoveAgent(actionBuffers.DiscreteActions);
+        if (!frozen)
+        {
+            // Move the agent using the action.
+            MoveAgent(actionBuffers.DiscreteActions);
+        }
 
         // Penalty given each step to encourage agent to finish task quickly.
         AddReward(-1f / MaxStep);

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

             {
                 return;
             }
-            item.Agent.EndEpisode();
         }
 
         //Disable killed Agent
             {
-                item.Col.gameObject.SetActive(false);
-                break;
+                // Disable collider
+                item.Agent.frozen = true;
+                item.Rb.constraints |= RigidbodyConstraints.FreezePositionY;
+                item.Col.enabled = false;
-        //End Episode
-        foreach (var item in ZombiesList)
-        {
-            if (item.Agent.transform == t)
-            {
-                item.Agent.gameObject.SetActive(false);
-                break;
-            }
-        }
+        // //End Episode
+        // foreach (var item in ZombiesList)
+        // {
+        //     if (item.Agent.transform == t)
+        //     {
+        //         item.Agent.gameObject.SetActive(false);
+        //         break;
+        //     }
+        // }
     }
 
 
             item.Agent.transform.SetPositionAndRotation(pos, rot);
             item.Rb.velocity = Vector3.zero;
             item.Rb.angularVelocity = Vector3.zero;
+            item.Col.enabled = true;
+            item.Agent.frozen = false;
+            item.Rb.constraints &= ~RigidbodyConstraints.FreezePositionY;
         }
 
         //Reset Blocks