[bug-fix] Fix stats reporting for reward signals in SAC (#3606)

Author: Ervin T

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

@@ -56,7 +56,6 @@ def __init__(self, seed, brain, trainer_parameters, load=False):
 
         self.use_recurrent = trainer_parameters["use_recurrent"]
         self.memory_dict: Dict[str, np.ndarray] = {}
-        self.reward_signals: Dict[str, "RewardSignal"] = {}
         self.num_branches = len(self.brain.vector_action_space_size)
         self.previous_action_dict: Dict[str, np.array] = {}
         self.normalize = trainer_parameters.get("normalize", False)

ml-agents/mlagents/trainers/sac/trainer.py

@@ -156,7 +156,7 @@ def _process_trajectory(self, trajectory: Trajectory) -> None:
         self.collected_rewards["environment"][agent_id] += np.sum(
             agent_buffer_trajectory["environment_rewards"]
         )
-        for name, reward_signal in self.policy.reward_signals.items():
+        for name, reward_signal in self.optimizer.reward_signals.items():
             evaluate_result = reward_signal.evaluate_batch(
                 agent_buffer_trajectory
             ).scaled_reward
@@ -223,9 +223,6 @@ def create_policy(self, brain_parameters: BrainParameters) -> TFPolicy:
             reparameterize=True,
             create_tf_graph=False,
         )
-        for _reward_signal in policy.reward_signals.keys():
-            self.collected_rewards[_reward_signal] = defaultdict(lambda: 0)
-
         # Load the replay buffer if load
         if self.load and self.checkpoint_replay_buffer:
             try:

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

@@ -221,6 +221,10 @@ def test_process_trajectory(dummy_config):
         for agent in reward.values():
             assert agent == 0
     assert trainer.stats_reporter.get_stats_summaries("Policy/Extrinsic Reward").num > 0
+    # Assert we're not just using the default values
+    assert (
+        trainer.stats_reporter.get_stats_summaries("Policy/Extrinsic Reward").mean > 0
+    )
 
 
 if __name__ == "__main__":