diff --git a/src/gfn/containers/replay_buffer_manager.py b/src/gfn/containers/replay_buffer_manager.py index 32e500f3..53465812 100644 --- a/src/gfn/containers/replay_buffer_manager.py +++ b/src/gfn/containers/replay_buffer_manager.py @@ -142,10 +142,31 @@ def _apply_baseline_sync(self, msg) -> None: Sent by a coordinator aggregating local baselines across buffer managers. The manager doesn't need to know who sent it or why. + + A malformed message is an error, not a no-op: the adopted value + gates what every assigned worker sends, so a manager that quietly + keeps its shard-local threshold filters against a different cutoff + from its siblings — the exact condition this message exists to + remove, and one that leaves no trace in the run. """ - global_baseline = msg.message_data.get("global_baseline_log_reward") - if global_baseline is not None: - self._global_baseline = float(global_baseline) + if not isinstance(msg.message_data, dict): + raise ValueError( + "BASELINE_SYNC message_data must be a dict, got " + f"{type(msg.message_data).__name__}" + ) + value = msg.message_data.get("global_baseline_log_reward") + if value is None: + raise ValueError( + "BASELINE_SYNC has no usable 'global_baseline_log_reward'; " + f"message_data={msg.message_data}" + ) + global_baseline = float(value) + if not math.isfinite(global_baseline): + raise ValueError( + f"BASELINE_SYNC global_baseline_log_reward is {global_baseline}; " + "a non-finite baseline keeps every trajectory or none" + ) + self._global_baseline = global_baseline def _inject_baseline_log_reward( self, score_dict: dict[str, float], incoming diff --git a/testing/test_replay_buffer.py b/testing/test_replay_buffer.py index 45d0b3fe..2fdefd69 100644 --- a/testing/test_replay_buffer.py +++ b/testing/test_replay_buffer.py @@ -790,7 +790,19 @@ def test_manager_baseline_sync_overrides_local(simple_env, trajectories): assert score["baseline_log_reward"] == pytest.approx(-7.5) -def test_manager_baseline_sync_ignores_missing_payload(simple_env, trajectories): +@pytest.mark.parametrize( + "payload", + [ + {}, + {"global_baseline_log_reward": None}, + {"global_baseline_log_reward": float("-inf")}, + None, + ], + ids=["missing", "none", "non_finite", "not_a_dict"], +) +def test_manager_baseline_sync_rejects_malformed_payload(simple_env, payload): + # Keeping the shard-local threshold on a malformed sync is the + # inconsistency this message exists to remove, and it leaves no trace. mgr = ReplayBufferManager( simple_env, rank=0, @@ -798,10 +810,9 @@ def test_manager_baseline_sync_ignores_missing_payload(simple_env, trajectories) capacity=5, store_locally=False, ) - mgr._apply_baseline_sync(Message(MessageType.BASELINE_SYNC, {})) - score: dict = {"score": 0.0} - mgr._inject_baseline_log_reward(score, trajectories) - assert score["baseline_log_reward"] == pytest.approx(0.0) # falls back to EMA. + with pytest.raises(ValueError, match="BASELINE_SYNC"): + mgr._apply_baseline_sync(Message(MessageType.BASELINE_SYNC, payload)) + assert mgr._global_baseline is None def test_manager_multi_manager_baseline_warning(simple_env):