cherry pick PR#3032 (#3066)

Author: Chris Elion

ml-agents/mlagents/trainers/ppo/policy.py

@@ -249,10 +249,13 @@ def get_value_estimates(
             ]
         if self.use_vec_obs:
             feed_dict[self.model.vector_in] = [brain_info.vector_observations[idx]]
+        agent_id = brain_info.agents[idx]
         if self.use_recurrent:
-            feed_dict[self.model.memory_in] = self.retrieve_memories([idx])
+            feed_dict[self.model.memory_in] = self.retrieve_memories([agent_id])
         if not self.use_continuous_act and self.use_recurrent:
-            feed_dict[self.model.prev_action] = self.retrieve_previous_action([idx])
+            feed_dict[self.model.prev_action] = self.retrieve_previous_action(
+                [agent_id]
+            )
         value_estimates = self.sess.run(self.model.value_heads, feed_dict)
 
         value_estimates = {k: float(v) for k, v in value_estimates.items()}

ml-agents/mlagents/trainers/tf_policy.py

@@ -56,9 +56,9 @@ def __init__(self, seed, brain, trainer_parameters):
         self.seed = seed
         self.brain = brain
         self.use_recurrent = trainer_parameters["use_recurrent"]
-        self.memory_dict: Dict[int, np.ndarray] = {}
+        self.memory_dict: Dict[str, np.ndarray] = {}
         self.num_branches = len(self.brain.vector_action_space_size)
-        self.previous_action_dict: Dict[int, np.array] = {}
+        self.previous_action_dict: Dict[str, np.array] = {}
         self.normalize = trainer_parameters.get("normalize", False)
         self.use_continuous_act = brain.vector_action_space_type == "continuous"
         if self.use_continuous_act:
@@ -181,14 +181,14 @@ def make_empty_memory(self, num_agents):
         return np.zeros((num_agents, self.m_size), dtype=np.float)
 
     def save_memories(
-        self, agent_ids: List[int], memory_matrix: Optional[np.ndarray]
+        self, agent_ids: List[str], memory_matrix: Optional[np.ndarray]
     ) -> None:
         if memory_matrix is None:
             return
         for index, agent_id in enumerate(agent_ids):
             self.memory_dict[agent_id] = memory_matrix[index, :]
 
-    def retrieve_memories(self, agent_ids: List[int]) -> np.ndarray:
+    def retrieve_memories(self, agent_ids: List[str]) -> np.ndarray:
         memory_matrix = np.zeros((len(agent_ids), self.m_size), dtype=np.float)
         for index, agent_id in enumerate(agent_ids):
             if agent_id in self.memory_dict:
@@ -209,14 +209,14 @@ def make_empty_previous_action(self, num_agents):
         return np.zeros((num_agents, self.num_branches), dtype=np.int)
 
     def save_previous_action(
-        self, agent_ids: List[int], action_matrix: Optional[np.ndarray]
+        self, agent_ids: List[str], action_matrix: Optional[np.ndarray]
     ) -> None:
         if action_matrix is None:
             return
         for index, agent_id in enumerate(agent_ids):
             self.previous_action_dict[agent_id] = action_matrix[index, :]
 
-    def retrieve_previous_action(self, agent_ids: List[int]) -> np.ndarray:
+    def retrieve_previous_action(self, agent_ids: List[str]) -> np.ndarray:
         action_matrix = np.zeros((len(agent_ids), self.num_branches), dtype=np.int)
         for index, agent_id in enumerate(agent_ids):
             if agent_id in self.previous_action_dict: