Replace BrainInfos with BatchedStepResult (#3207)

Author: Ervin T

gym-unity/gym_unity/envs/__init__.py

@@ -177,7 +177,7 @@ def step(self, action: List[Any]) -> GymStepResult:
             observation (object/list): agent's observation of the current environment
             reward (float/list) : amount of reward returned after previous action
             done (boolean/list): whether the episode has ended.
-            info (dict): contains auxiliary diagnostic information, including BrainInfo.
+            info (dict): contains auxiliary diagnostic information, including BatchedStepResult.
         """
 
         # Use random actions for all other agents in environment.

gym-unity/gym_unity/tests/test_gym.py

@@ -9,9 +9,9 @@
 
 @mock.patch("gym_unity.envs.UnityEnvironment")
 def test_gym_wrapper(mock_env):
-    mock_brain = create_mock_group_spec()
-    mock_braininfo = create_mock_vector_step_result()
-    setup_mock_unityenvironment(mock_env, mock_brain, mock_braininfo)
+    mock_spec = create_mock_group_spec()
+    mock_step = create_mock_vector_step_result()
+    setup_mock_unityenvironment(mock_env, mock_spec, mock_step)
 
     env = UnityEnv(" ", use_visual=False, multiagent=False)
     assert isinstance(env, UnityEnv)
@@ -28,9 +28,9 @@ def test_gym_wrapper(mock_env):
 
 @mock.patch("gym_unity.envs.UnityEnvironment")
 def test_multi_agent(mock_env):
-    mock_brain = create_mock_group_spec()
-    mock_braininfo = create_mock_vector_step_result(num_agents=2)
-    setup_mock_unityenvironment(mock_env, mock_brain, mock_braininfo)
+    mock_spec = create_mock_group_spec()
+    mock_step = create_mock_vector_step_result(num_agents=2)
+    setup_mock_unityenvironment(mock_env, mock_spec, mock_step)
 
     with pytest.raises(UnityGymException):
         UnityEnv(" ", multiagent=False)
@@ -47,11 +47,11 @@ def test_multi_agent(mock_env):
 
 @mock.patch("gym_unity.envs.UnityEnvironment")
 def test_branched_flatten(mock_env):
-    mock_brain = create_mock_group_spec(
+    mock_spec = create_mock_group_spec(
         vector_action_space_type="discrete", vector_action_space_size=[2, 2, 3]
     )
-    mock_braininfo = create_mock_vector_step_result(num_agents=1)
-    setup_mock_unityenvironment(mock_env, mock_brain, mock_braininfo)
+    mock_step = create_mock_vector_step_result(num_agents=1)
+    setup_mock_unityenvironment(mock_env, mock_spec, mock_step)
 
     env = UnityEnv(" ", use_visual=False, multiagent=False, flatten_branched=True)
     assert isinstance(env.action_space, spaces.Discrete)
@@ -67,9 +67,9 @@ def test_branched_flatten(mock_env):
 @pytest.mark.parametrize("use_uint8", [True, False], ids=["float", "uint8"])
 @mock.patch("gym_unity.envs.UnityEnvironment")
 def test_gym_wrapper_visual(mock_env, use_uint8):
-    mock_brain = create_mock_group_spec(number_visual_observations=1)
-    mock_braininfo = create_mock_vector_step_result(number_visual_observations=1)
-    setup_mock_unityenvironment(mock_env, mock_brain, mock_braininfo)
+    mock_spec = create_mock_group_spec(number_visual_observations=1)
+    mock_step = create_mock_vector_step_result(number_visual_observations=1)
+    setup_mock_unityenvironment(mock_env, mock_spec, mock_step)
 
     env = UnityEnv(" ", use_visual=True, multiagent=False, uint8_visual=use_uint8)
     assert isinstance(env, UnityEnv)
@@ -117,10 +117,10 @@ def create_mock_group_spec(
 
 def create_mock_vector_step_result(num_agents=1, number_visual_observations=0):
     """
-    Creates a mock BrainInfo with vector observations. Imitates constant
+    Creates a mock BatchedStepResult with vector observations. Imitates constant
     vector observations, rewards, dones, and agents.
 
-    :int num_agents: Number of "agents" to imitate in your BrainInfo values.
+    :int num_agents: Number of "agents" to imitate in your BatchedStepResult values.
     """
     obs = [np.array([num_agents * [1, 2, 3]])]
     if number_visual_observations:
@@ -134,7 +134,7 @@ def create_mock_vector_step_result(num_agents=1, number_visual_observations=0):
 def setup_mock_unityenvironment(mock_env, mock_spec, mock_result):
     """
     Takes a mock UnityEnvironment and adds the appropriate properties, defined by the mock
-    BrainParameters and BrainInfo.
+    GroupSpec and BatchedStepResult.
 
     :Mock mock_env: A mock UnityEnvironment, usually empty.
     :Mock mock_spec: An AgentGroupSpec object that specifies the params of this environment.

ml-agents-envs/mlagents_envs/base_env.py

@@ -90,14 +90,24 @@ def __init__(self, obs, reward, done, max_step, agent_id, action_mask):
         self.max_step: np.ndarray = max_step
         self.agent_id: np.ndarray = agent_id
         self.action_mask: Optional[List[np.ndarray]] = action_mask
-        self._agent_id_to_index: Optional[Dict[int, int]] = None
+        self._agent_id_to_index: Optional[Dict[AgentId, int]] = None
 
-    def contains_agent(self, agent_id: AgentId) -> bool:
+    @property
+    def agent_id_to_index(self) -> Dict[AgentId, int]:
+        """
+        Returns the index of the agent_id in this BatchedStepResult, and
+        -1 if agent_id is not in this BatchedStepResult.
+        :param agent_id: The id of the agent
+        :returns: The index of the agent_id, and -1 if not found.
+        """
         if self._agent_id_to_index is None:
             self._agent_id_to_index = {}
             for a_idx, a_id in enumerate(self.agent_id):
                 self._agent_id_to_index[a_id] = a_idx
-        return agent_id in self._agent_id_to_index
+        return self._agent_id_to_index
+
+    def contains_agent(self, agent_id: AgentId) -> bool:
+        return agent_id in self.agent_id_to_index
 
     def get_agent_step_result(self, agent_id: AgentId) -> StepResult:
         """

ml-agents/mlagents/trainers/action_info.py

@@ -1,5 +1,6 @@
 from typing import NamedTuple, Any, Dict, List
 import numpy as np
+from mlagents_envs.base_env import AgentId
 
 ActionInfoOutputs = Dict[str, np.ndarray]
 
@@ -8,4 +9,4 @@ class ActionInfo(NamedTuple):
     action: Any
     value: Any
     outputs: ActionInfoOutputs
-    agents: List[str]
+    agent_ids: List[AgentId]

ml-agents/mlagents/trainers/agent_processor.py

@@ -1,13 +1,15 @@
 import sys
+import numpy as np
 from typing import List, Dict, Deque, TypeVar, Generic
 from collections import defaultdict, Counter, deque
 
+from mlagents_envs.base_env import BatchedStepResult
 from mlagents.trainers.trajectory import Trajectory, AgentExperience
-from mlagents.trainers.brain import BrainInfo
 from mlagents.trainers.tf_policy import TFPolicy
 from mlagents.trainers.policy import Policy
 from mlagents.trainers.action_info import ActionInfo, ActionInfoOutputs
 from mlagents.trainers.stats import StatsReporter
+from mlagents.trainers.env_manager import get_global_agent_id
 
 T = TypeVar("T")
 
@@ -35,7 +37,7 @@ def __init__(
         :param stats_category: The category under which to write the stats. Usually, this comes from the Trainer.
         """
         self.experience_buffers: Dict[str, List[AgentExperience]] = defaultdict(list)
-        self.last_brain_info: Dict[str, BrainInfo] = {}
+        self.last_step_result: Dict[str, BatchedStepResult] = {}
         # last_take_action_outputs stores the action a_t taken before the current observation s_(t+1), while
         # grabbing previous_action from the policy grabs the action PRIOR to that, a_(t-1).
         self.last_take_action_outputs: Dict[str, ActionInfoOutputs] = {}
@@ -50,12 +52,15 @@ def __init__(
         self.behavior_id = behavior_id
 
     def add_experiences(
-        self, curr_info: BrainInfo, previous_action: ActionInfo
+        self,
+        batched_step_result: BatchedStepResult,
+        worker_id: int,
+        previous_action: ActionInfo,
     ) -> None:
         """
         Adds experiences to each agent's experience history.
-        :param curr_info: current BrainInfo.
-        :param previous_action: The return value of the Policy's get_action method.
+        :param batched_step_result: current BatchedStepResult.
+        :param previous_action: The outputs of the Policy's get_action method.
         """
         take_action_outputs = previous_action.outputs
         if take_action_outputs:
@@ -65,99 +70,101 @@ def add_experiences(
                 "Policy/Learning Rate", take_action_outputs["learning_rate"]
             )
 
-        for agent_id in previous_action.agents:
-            self.last_take_action_outputs[agent_id] = take_action_outputs
-
-        # Store the environment reward
-        tmp_environment_reward = curr_info.rewards
-
-        for agent_idx, agent_id in enumerate(curr_info.agents):
-            stored_info = self.last_brain_info.get(agent_id, None)
+        # Make unique agent_ids that are global across workers
+        action_global_agent_ids = [
+            get_global_agent_id(worker_id, ag_id) for ag_id in previous_action.agent_ids
+        ]
+        for global_id in action_global_agent_ids:
+            self.last_take_action_outputs[global_id] = take_action_outputs
+
+        for _id in np.nditer(batched_step_result.agent_id):  # Explicit numpy iteration
+            local_id = int(
+                _id
+            )  # Needed for mypy to pass since ndarray has no content type
+            curr_agent_step = batched_step_result.get_agent_step_result(local_id)
+            global_id = get_global_agent_id(worker_id, local_id)
+            stored_step = self.last_step_result.get(global_id, None)
             stored_take_action_outputs = self.last_take_action_outputs.get(
-                agent_id, None
+                global_id, None
             )
-            if stored_info is not None and stored_take_action_outputs is not None:
-                prev_idx = stored_info.agents.index(agent_id)
-                obs = []
-                if not stored_info.local_done[prev_idx]:
-                    for i, _ in enumerate(stored_info.visual_observations):
-                        obs.append(stored_info.visual_observations[i][prev_idx])
-                    if self.policy.use_vec_obs:
-                        obs.append(stored_info.vector_observations[prev_idx])
+            if stored_step is not None and stored_take_action_outputs is not None:
+                # We know the step is from the same worker, so use the local agent id.
+                stored_agent_step = stored_step.get_agent_step_result(local_id)
+                idx = stored_step.agent_id_to_index[local_id]
+                obs = stored_agent_step.obs
+                if not stored_agent_step.done:
                     if self.policy.use_recurrent:
-                        memory = self.policy.retrieve_memories([agent_id])[0, :]
+                        memory = self.policy.retrieve_memories([global_id])[0, :]
                     else:
                         memory = None
 
-                    done = curr_info.local_done[agent_idx]
-                    max_step = curr_info.max_reached[agent_idx]
+                    done = curr_agent_step.done
+                    max_step = curr_agent_step.max_step
 
                     # Add the outputs of the last eval
-                    action = stored_take_action_outputs["action"][prev_idx]
+                    action = stored_take_action_outputs["action"][idx]
                     if self.policy.use_continuous_act:
-                        action_pre = stored_take_action_outputs["pre_action"][prev_idx]
+                        action_pre = stored_take_action_outputs["pre_action"][idx]
                     else:
                         action_pre = None
-                    action_probs = stored_take_action_outputs["log_probs"][prev_idx]
-                    action_masks = stored_info.action_masks[prev_idx]
-                    prev_action = self.policy.retrieve_previous_action([agent_id])[0, :]
+                    action_probs = stored_take_action_outputs["log_probs"][idx]
+                    action_mask = stored_agent_step.action_mask
+                    prev_action = self.policy.retrieve_previous_action([global_id])[
+                        0, :
+                    ]
 
                     experience = AgentExperience(
                         obs=obs,
-                        reward=tmp_environment_reward[agent_idx],
+                        reward=curr_agent_step.reward,
                         done=done,
                         action=action,
                         action_probs=action_probs,
                         action_pre=action_pre,
-                        action_mask=action_masks,
+                        action_mask=action_mask,
                         prev_action=prev_action,
                         max_step=max_step,
                         memory=memory,
                     )
                     # Add the value outputs if needed
-                    self.experience_buffers[agent_id].append(experience)
-                    self.episode_rewards[agent_id] += tmp_environment_reward[agent_idx]
+                    self.experience_buffers[global_id].append(experience)
+                    self.episode_rewards[global_id] += curr_agent_step.reward
                 if (
-                    curr_info.local_done[agent_idx]
+                    curr_agent_step.done
                     or (
-                        len(self.experience_buffers[agent_id])
+                        len(self.experience_buffers[global_id])
                         >= self.max_trajectory_length
                     )
-                ) and len(self.experience_buffers[agent_id]) > 0:
+                ) and len(self.experience_buffers[global_id]) > 0:
                     # Make next AgentExperience
-                    next_obs = []
-                    for i, _ in enumerate(curr_info.visual_observations):
-                        next_obs.append(curr_info.visual_observations[i][agent_idx])
-                    if self.policy.use_vec_obs:
-                        next_obs.append(curr_info.vector_observations[agent_idx])
+                    next_obs = curr_agent_step.obs
                     trajectory = Trajectory(
-                        steps=self.experience_buffers[agent_id],
-                        agent_id=agent_id,
+                        steps=self.experience_buffers[global_id],
+                        agent_id=global_id,
                         next_obs=next_obs,
                         behavior_id=self.behavior_id,
                     )
                     for traj_queue in self.trajectory_queues:
                         traj_queue.put(trajectory)
-                    self.experience_buffers[agent_id] = []
-                    if curr_info.local_done[agent_idx]:
+                    self.experience_buffers[global_id] = []
+                    if curr_agent_step.done:
                         self.stats_reporter.add_stat(
                             "Environment/Cumulative Reward",
-                            self.episode_rewards.get(agent_id, 0),
+                            self.episode_rewards.get(global_id, 0),
                         )
                         self.stats_reporter.add_stat(
                             "Environment/Episode Length",
-                            self.episode_steps.get(agent_id, 0),
+                            self.episode_steps.get(global_id, 0),
                         )
-                        del self.episode_steps[agent_id]
-                        del self.episode_rewards[agent_id]
-                elif not curr_info.local_done[agent_idx]:
-                    self.episode_steps[agent_id] += 1
+                        del self.episode_steps[global_id]
+                        del self.episode_rewards[global_id]
+                elif not curr_agent_step.done:
+                    self.episode_steps[global_id] += 1
 
-            self.last_brain_info[agent_id] = curr_info
+            self.last_step_result[global_id] = batched_step_result
 
         if "action" in take_action_outputs:
             self.policy.save_previous_action(
-                previous_action.agents, take_action_outputs["action"]
+                previous_action.agent_ids, take_action_outputs["action"]
             )
 
     def publish_trajectory_queue(