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

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

 
         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

         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
             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

         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__":