Fix some of the tests

ml-agents/mlagents/trainers/agent_processor.py

 
                 if (
                     next_info.local_done[next_idx]
-                    or len(self.experience_buffers[agent_id]) > self.time_horizon
+                    or len(self.experience_buffers[agent_id]) >= self.time_horizon
                 ) and len(self.experience_buffers[agent_id]) > 0:
                     # Make next AgentExperience
                     next_obs = []

ml-agents/mlagents/trainers/rl_trainer.py

 # # Unity ML-Agents Toolkit
 import logging
-from typing import Dict, List, Any, NamedTuple
+from typing import Dict, NamedTuple
-from mlagents.envs.brain import BrainInfo
 from mlagents.envs.action_info import ActionInfoOutputs
 from mlagents.trainers.buffer import AgentBuffer
 from mlagents.trainers.trainer import Trainer, UnityTrainerException
         self.update_buffer = AgentBuffer()
         self.episode_steps = {}
 
-    def construct_curr_info(self, next_info: BrainInfo) -> BrainInfo:
-        """
-        Constructs a BrainInfo which contains the most recent previous experiences for all agents
-        which correspond to the agents in a provided next_info.
-        :BrainInfo next_info: A t+1 BrainInfo.
-        :return: curr_info: Reconstructed BrainInfo to match agents of next_info.
-        """
-        visual_observations: List[List[Any]] = [
-            [] for _ in next_info.visual_observations
-        ]  # TODO add types to brain.py methods
-        vector_observations = []
-        rewards = []
-        local_dones = []
-        max_reacheds = []
-        agents = []
-        action_masks = []
-        for agent_id in next_info.agents:
-            agent_brain_info = self.processing_buffer[agent_id].last_brain_info
-            if agent_brain_info is None:
-                agent_brain_info = next_info
-            agent_index = agent_brain_info.agents.index(agent_id)
-            for i in range(len(next_info.visual_observations)):
-                visual_observations[i].append(
-                    agent_brain_info.visual_observations[i][agent_index]
-                )
-            vector_observations.append(
-                agent_brain_info.vector_observations[agent_index]
-            )
-            rewards.append(agent_brain_info.rewards[agent_index])
-            local_dones.append(agent_brain_info.local_done[agent_index])
-            max_reacheds.append(agent_brain_info.max_reached[agent_index])
-            agents.append(agent_brain_info.agents[agent_index])
-            action_masks.append(agent_brain_info.action_masks[agent_index])
-        curr_info = BrainInfo(
-            visual_observations,
-            vector_observations,
-            rewards,
-            agents,
-            local_dones,
-            max_reacheds,
-            action_masks,
-        )
-        return curr_info
-
-    def add_experiences(
-        self,
-        curr_info: BrainInfo,
-        next_info: BrainInfo,
-        take_action_outputs: ActionInfoOutputs,
-    ) -> None:
-        """
-        Adds experiences to each agent's experience history.
-        :param curr_info: current BrainInfo.
-        :param next_info: next BrainInfo.
-        :param take_action_outputs: The outputs of the Policy's get_action method.
-        """
-        self.trainer_metrics.start_experience_collection_timer()
-        if take_action_outputs:
-            self.stats["Policy/Entropy"].append(take_action_outputs["entropy"].mean())
-            self.stats["Policy/Learning Rate"].append(
-                take_action_outputs["learning_rate"]
-            )
-            for name, signal in self.policy.reward_signals.items():
-                self.stats[signal.value_name].append(
-                    np.mean(take_action_outputs["value_heads"][name])
-                )
-
-        for agent_id in curr_info.agents:
-            self.processing_buffer[agent_id].last_brain_info = curr_info
-            self.processing_buffer[
-                agent_id
-            ].last_take_action_outputs = take_action_outputs
-
-        if curr_info.agents != next_info.agents:
-            curr_to_use = self.construct_curr_info(next_info)
-        else:
-            curr_to_use = curr_info
-
-        # Evaluate and store the reward signals
-        tmp_reward_signal_outs = {}
-        for name, signal in self.policy.reward_signals.items():
-            tmp_reward_signal_outs[name] = signal.evaluate(
-                curr_to_use, take_action_outputs["action"], next_info
-            )
-        # Store the environment reward
-        tmp_environment = np.array(next_info.rewards, dtype=np.float32)
-
-        rewards_out = AllRewardsOutput(
-            reward_signals=tmp_reward_signal_outs, environment=tmp_environment
-        )
-
-        for agent_id in next_info.agents:
-            stored_info = self.processing_buffer[agent_id].last_brain_info
-            stored_take_action_outputs = self.processing_buffer[
-                agent_id
-            ].last_take_action_outputs
-            if stored_info is not None:
-                idx = stored_info.agents.index(agent_id)
-                next_idx = next_info.agents.index(agent_id)
-                if not stored_info.local_done[idx]:
-                    for i, _ in enumerate(stored_info.visual_observations):
-                        self.processing_buffer[agent_id]["visual_obs%d" % i].append(
-                            stored_info.visual_observations[i][idx]
-                        )
-                        self.processing_buffer[agent_id][
-                            "next_visual_obs%d" % i
-                        ].append(next_info.visual_observations[i][next_idx])
-                    if self.policy.use_vec_obs:
-                        self.processing_buffer[agent_id]["vector_obs"].append(
-                            stored_info.vector_observations[idx]
-                        )
-                        self.processing_buffer[agent_id]["next_vector_in"].append(
-                            next_info.vector_observations[next_idx]
-                        )
-                    if self.policy.use_recurrent:
-                        self.processing_buffer[agent_id]["memory"].append(
-                            self.policy.retrieve_memories([agent_id])[0, :]
-                        )
-
-                    self.processing_buffer[agent_id]["masks"].append(1.0)
-                    self.processing_buffer[agent_id]["done"].append(
-                        next_info.local_done[next_idx]
-                    )
-                    # Add the outputs of the last eval
-                    self.add_policy_outputs(stored_take_action_outputs, agent_id, idx)
-                    # Store action masks if necessary
-                    if not self.policy.use_continuous_act:
-                        self.processing_buffer[agent_id]["action_mask"].append(
-                            stored_info.action_masks[idx], padding_value=1
-                        )
-                    self.processing_buffer[agent_id]["prev_action"].append(
-                        self.policy.retrieve_previous_action([agent_id])[0, :]
-                    )
-
-                    values = stored_take_action_outputs["value_heads"]
-
-                    # Add the value outputs if needed
-                    self.add_rewards_outputs(
-                        rewards_out, values, agent_id, idx, next_idx
-                    )
-
-                    for name, rewards in self.collected_rewards.items():
-                        if agent_id not in rewards:
-                            rewards[agent_id] = 0
-                        if name == "environment":
-                            # Report the reward from the environment
-                            rewards[agent_id] += rewards_out.environment[next_idx]
-                        else:
-                            # Report the reward signals
-                            rewards[agent_id] += rewards_out.reward_signals[
-                                name
-                            ].scaled_reward[next_idx]
-                if not next_info.local_done[next_idx]:
-                    if agent_id not in self.episode_steps:
-                        self.episode_steps[agent_id] = 0
-                    self.episode_steps[agent_id] += 1
-        self.policy.save_previous_action(
-            curr_info.agents, take_action_outputs["action"]
-        )
-        self.trainer_metrics.end_experience_collection_timer()
-
-        self.processing_buffer.reset_local_buffers()
         for agent_id in self.episode_steps:
             self.episode_steps[agent_id] = 0
         for rewards in self.collected_rewards.values():

ml-agents/mlagents/trainers/tests/test_rl_trainer.py

     return mock_policy
 
 
-@mock.patch("mlagents.trainers.rl_trainer.RLTrainer.add_policy_outputs")
-@mock.patch("mlagents.trainers.rl_trainer.RLTrainer.add_rewards_outputs")
-@pytest.mark.parametrize("num_vis_obs", [0, 1, 2], ids=["vec", "1 viz", "2 viz"])
-def test_rl_trainer(add_policy_outputs, add_rewards_outputs, num_vis_obs):
+def test_rl_trainer():
-    trainer.policy = create_mock_policy()
-    fake_action_outputs = {
-        "action": [0.1, 0.1],
-        "value_heads": {},
-        "entropy": np.array([1.0], dtype=np.float32),
-        "learning_rate": 1.0,
-    }
-    mock_braininfo = mb.create_mock_braininfo(
-        num_agents=2,
-        num_vector_observations=8,
-        num_vector_acts=2,
-        num_vis_observations=num_vis_obs,
-    )
-    trainer.add_experiences(mock_braininfo, mock_braininfo, fake_action_outputs)
-
-    # Remove one of the agents
-    next_mock_braininfo = mb.create_mock_braininfo(
-        num_agents=1,
-        num_vector_observations=8,
-        num_vector_acts=2,
-        num_vis_observations=num_vis_obs,
-    )
-    brain_info = trainer.construct_curr_info(next_mock_braininfo)
-
-    # assert construct_curr_info worked properly
-    assert len(brain_info.agents) == 1
-    assert len(brain_info.visual_observations) == num_vis_obs
-    assert len(brain_info.vector_observations) == 1
-
+    agent_id = "0"
+    trainer.episode_steps[agent_id] = 3
+    trainer.collected_rewards["extrinsic"] = {agent_id: 3}
-        assert len(trainer.processing_buffer[agent_id]["action"]) == 0
     for rewards in trainer.collected_rewards.values():
         for agent_id in rewards:
             assert rewards[agent_id] == 0

ml-agents/mlagents/trainers/tests/test_agent_processor.py

+import unittest.mock as mock
+import pytest
+import mlagents.trainers.tests.mock_brain as mb
+import numpy as np
+from mlagents.trainers.agent_processor import AgentProcessor
+
+
+def create_mock_brain():
+    mock_brain = mb.create_mock_brainparams(
+        vector_action_space_type="continuous",
+        vector_action_space_size=[2],
+        vector_observation_space_size=8,
+        number_visual_observations=1,
+    )
+    return mock_brain
+
+
+def create_mock_policy():
+    mock_policy = mock.Mock()
+    mock_policy.reward_signals = {}
+    mock_policy.retrieve_memories.return_value = np.zeros((1, 1), dtype=np.float32)
+    mock_policy.retrieve_previous_action.return_value = np.zeros(
+        (1, 1), dtype=np.float32
+    )
+    return mock_policy
+
+
+@mock.patch("mlagents.trainers.rl_trainer.RLTrainer.add_policy_outputs")
+@mock.patch("mlagents.trainers.rl_trainer.RLTrainer.add_rewards_outputs")
+@pytest.mark.parametrize("num_vis_obs", [0, 1, 2], ids=["vec", "1 viz", "2 viz"])
+def test_agentprocessor(add_policy_outputs, add_rewards_outputs, num_vis_obs):
+    policy = create_mock_policy()
+    trainer = mock.Mock()
+    processor = AgentProcessor(trainer, policy, time_horizon=5)
+    fake_action_outputs = {
+        "action": [0.1, 0.1],
+        "value_heads": {},
+        "entropy": np.array([1.0], dtype=np.float32),
+        "learning_rate": 1.0,
+        "pre_action": [0.1, 0.1],
+        "log_probs": [0.1, 0.1],
+    }
+    mock_braininfo = mb.create_mock_braininfo(
+        num_agents=2,
+        num_vector_observations=8,
+        num_vector_acts=2,
+        num_vis_observations=num_vis_obs,
+    )
+    for i in range(5):
+        processor.add_experiences(mock_braininfo, mock_braininfo, fake_action_outputs)
+
+    # Assert that two trajectories have been added to the Trainer
+    assert len(trainer.process_trajectory.call_args_list) == 2
+
+    # Assert that the trajectory is of length 5
+    trajectory = trainer.process_trajectory.call_args_list[0][0][0]
+    assert len(trajectory.steps) == 5
+
+    # Assert that the AgentProcessor is empty
+    assert len(processor.experience_buffers[0]) == 0