fix: action type

pettingzoo/utils/conversions.py

@@ -1,11 +1,16 @@
 import copy
 import warnings
 from collections import defaultdict
+from typing import Dict, Optional
 
 from pettingzoo.utils import agent_selector
 from pettingzoo.utils.env import AECEnv, ParallelEnv
 from pettingzoo.utils.wrappers import OrderEnforcingWrapper
 
+ActionType = Optional[int]
+AgentID = str
+ActionDict = Dict[AgentID, ActionType]
+
 
 def parallel_wrapper_fn(env_fn):
     def par_fn(**kwargs):
@@ -231,7 +236,7 @@ def reset(self, seed=None, return_info=False, options=None):
         self._observations = self.env.reset(seed=seed, options=options)
         self.agents = self.env.agents[:]
         self._live_agents = self.agents[:]
-        self._actions = {agent: None for agent in self.agents}
+        self._actions: ActionDict = {agent: None for agent in self.agents}
         self._agent_selector = agent_selector(self._live_agents)
         self.agent_selection = self._agent_selector.reset()
         self.dones = {agent: False for agent in self.agents}
@@ -257,9 +262,12 @@ def add_new_agent(self, new_agent):
         self.rewards[new_agent] = 0
         self._cumulative_rewards[new_agent] = 0
 
-    def step(self, action: None):
+    def step(self, action: ActionType):
+        if action is not None:
+            action = int(action)
         if self.dones[self.agent_selection]:
             del self._actions[self.agent_selection]
+            assert action is None
             return self._was_done_step(action)
         self._actions[self.agent_selection] = action
         if self._agent_selector.is_last():