From 7264420e9028e761acc41736f6f61e2e0e8e88ec Mon Sep 17 00:00:00 2001
From: Henry Peteet <henry.peteet@unity3d.com>
Date: Wed, 10 Nov 2021 10:04:07 -0500
Subject: [PATCH] Two small bugfixes (#5589)

1. Add the missing `_cumulative_rewards` property
2. Update `agent_selection` to not error out when an agent finishes an episode.
---
 pettingzoo-unity/pettingzoo_unity/pettingzoo_wrapper.py | 7 +++++++
 1 file changed, 7 insertions(+)

diff --git a/pettingzoo-unity/pettingzoo_unity/pettingzoo_wrapper.py b/pettingzoo-unity/pettingzoo_unity/pettingzoo_wrapper.py
index b306718853..d4238c3216 100644
--- a/pettingzoo-unity/pettingzoo_unity/pettingzoo_wrapper.py
+++ b/pettingzoo-unity/pettingzoo_unity/pettingzoo_wrapper.py
@@ -307,6 +307,10 @@ def _create_empty_actions(self, behavior_name, num_agents):
             np.zeros((num_agents, len(a_spec.discrete_branches)), dtype=np.int32),
         )
 
+    @property
+    def _cumulative_rewards(self):
+        return self._cumm_rewards
+
     def reset(self):
         """
         Resets the environment.
@@ -399,6 +403,9 @@ def infos(self):
 
     @property
     def agent_selection(self):
+        if not self._live_agents:
+            # If we had an agent finish then return that agent even though it isn't alive.
+            return self._agents[0]
         return self._agents[self._agent_index]
 
     @property