f

src/bsk_rl/data/composition.py

@@ -98,6 +98,10 @@ def update_with_communicated_data(self) -> None:
 class ComposedReward(GlobalReward):
     datastore_type = ComposedDataStore
 
+    def pass_data(self) -> Data:
+        for rewarder, data in zip(self.rewarders, self.data.data):
+            rewarder.data = data
+
     def __init__(self, *rewarders: GlobalReward) -> None:
         """Construct composed reward.
 
@@ -169,3 +173,8 @@ def calculate_reward(
                 for sat_id, sat_reward in reward_i.items():
                     reward[sat_id] = reward.get(sat_id, 0.0) + sat_reward
         return reward
+
+    def reward(self, new_data_dict: dict[str, ComposedData]) -> dict[str, float]:
+        reward = super().reward(new_data_dict)
+        self.pass_data()
+        return reward