Add baseline filtering and scoring options to ReplayBuffer#524
Merged
Conversation
- Introduced baseline_filtering to send only trajectories with log-reward above the baseline. - Added scoring_only to send only terminating states and log-rewards. - Updated methods to handle filtering and baseline updates accordingly. Co-authored-by: Copilot <copilot@github.com>
Codecov Report❌ Patch coverage is
Additional details and impacted files@@ Coverage Diff @@
## master #524 +/- ##
==========================================
+ Coverage 73.92% 74.39% +0.47%
==========================================
Files 59 59
Lines 10539 10636 +97
Branches 1523 1546 +23
==========================================
+ Hits 7791 7913 +122
+ Misses 2252 2219 -33
- Partials 496 504 +8
🚀 New features to boost your workflow:
|
younik
reviewed
May 19, 2026
younik
left a comment
younik
left a comment
Collaborator
There was a problem hiding this comment.
Thanks for this!
My only concern is that this code is a bit overfitted to our usecase, but we are merging to broader torchgfn.
I would land these changes anyway, but we need at least a full test coverage for all the Container types.
Comment on lines
+345
to
+355
| terminating_states = container.terminating_states | ||
| log_rewards = container.log_rewards | ||
| n = len(terminating_states) | ||
| return StatesContainer( | ||
| env=self.env, | ||
| states=terminating_states, | ||
| is_terminating=torch.ones( | ||
| n, dtype=torch.bool, device=terminating_states.device | ||
| ), | ||
| log_rewards=log_rewards, | ||
| ) |
Collaborator
There was a problem hiding this comment.
I suspect this doesn't work with full Trajectories. The script we are running are using terminal states as container, thus log_rewards refer to them, but in general, if container are trajectories, container.terminating_states and container.log_rewards have different shapes
Comment on lines
+308
to
+328
| log_rewards = container.log_rewards | ||
| if log_rewards is None: | ||
| return container # Can't filter without rewards. | ||
|
|
||
| mask = log_rewards >= self._baseline_log_reward | ||
|
|
||
| n_total = len(container) | ||
| n_kept = int(mask.sum().item()) | ||
| if self.timing: | ||
| self._baseline_total += n_total | ||
| self._baseline_kept += n_kept | ||
|
|
||
| if mask.all(): | ||
| return container # All above baseline — send everything. | ||
| if not mask.any(): | ||
| if self.timing: | ||
| self._baseline_skipped_sends += 1 | ||
| return None # All below baseline — skip send. | ||
|
|
||
| indices = torch.where(mask)[0] | ||
| return container[indices] |
Collaborator
There was a problem hiding this comment.
if we want to have this for general-use, we should keep the container type consistent.
Think about what happen with the current code if container is a trajectory, it seems to me we select only the states with reward higher to baseline, but I believe we should return the full trajectories based on a baseline constraint (e.g. sum of rewards in the trajectory).
Maybe container[indices] even crashes for trajectories because bad indexing, not sure
Address review feedback on the baseline-filtering and scoring-only paths:
Worker (src/gfn/containers/replay_buffer.py)
- Fix mangled docstring: separate the lazy_sort / baseline_filtering /
scoring_only entries and put the closing triple-quote on its own line.
- Reject Transitions in _filter_by_baseline and _prepare_for_remote:
Transitions.log_rewards is per-row with -inf for non-terminating
steps, so per-row masking would silently drop intermediate transitions
and break DB/SubTB. Document Trajectories / StatesContainer as the
supported payloads and assert at the entry points.
- Assert log_rewards shape and reject backward Trajectories in
_prepare_for_remote (no log_rewards => cannot build the lean
StatesContainer payload).
- Cross-validate at init: baseline_filtering=True now raises if
remote_manager_rank is None.
- Track filter counters unconditionally (no longer gated on timing=True).
- Add baseline_refresh_after (default 10): after this many consecutive
fully-filtered batches the worker bypasses the filter on the next send
so it can receive a fresh baseline. Prevents the worker from
permanently starving the manager once the baseline drifts above the
worker's current support.
- Forward the three new flags through
NormBasedDiversePrioritizedReplayBuffer.__init__.
Manager (src/gfn/containers/replay_buffer_manager.py)
- Add class docstring documenting store_locally and the baseline
injection contract.
- New args: baseline_strategy ("min" | "percentile" | "ema", default
"min"), baseline_percentile=0.1, baseline_ema_alpha=0.1. Strategy is
validated at init.
- _inject_baseline_log_reward now takes the incoming container and
maintains an EMA of incoming batch minima (with torch.isfinite
filtering so Transitions' -inf entries do not poison the statistic).
Dispatch:
- store_locally=True and buffer at capacity => min or percentile of
the buffer (per strategy).
- otherwise => EMA fallback.
This makes store_locally=False + baseline_filtering=True compose
correctly; previously the combination was a silent no-op because the
manager's empty buffer never tripped the len >= capacity gate.
Tests (testing/test_replay_buffer.py)
- 15 new tests covering: init-time validation, _filter_by_baseline
paths (no-baseline-yet / all-above / all-below / partial /
force-refresh-after-N-empties / rejects Transitions),
_prepare_for_remote paths (Trajectories -> StatesContainer /
StatesContainer no-op / rejects Transitions), and manager-side
baseline injection under min / percentile / EMA-fallback /
non-finite-skip strategies.
Pre-commit (black, flake8, pyright) clean; full pytest passes
(1927 passed, 735 skipped; no regressions).
Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com>
Collaborator
Author
|
Thank you Joseph for implementing these changes; the updates look great and really enhance the feature. Good to go on my end. |
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
3 participants
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Description
TLDR
More Details
This pull request introduces two major new features to the replay buffer system: baseline filtering and scoring-only communication. These changes allow agents to avoid sending low-value trajectories and to reduce communication costs by sending only essential data when appropriate. The buffer manager now also supports operation without local storage, and tracks/reporting for these new modes has been added.
Replay Buffer Improvements:
baseline_filteringoption toReplayBufferto filter out trajectories with log-reward below the current baseline (worst item in the remote buffer); includes logic for updating and applying the baseline, and tracking statistics on filtered/skipped sends. [1] [2] [3] [4] [5] [6] [7]scoring_onlyoption toReplayBuffer, allowing agents to send only terminating states and log-rewards (not full trajectories) to the remote buffer manager, reducing communication and serialization overhead. [1] [2] [3] [4]Replay Buffer Manager Enhancements:
store_locallyoption toReplayBufferManagerto allow operation without storing data locally; data is only scored and not retained if this is set toFalse. [1] [2] [3]baseline_log_rewardfield into score responses (when full) to communicate the current baseline to agents using baseline filtering. [1] [2] [3]Monitoring and Logging:
ReplayBufferto report statistics on baseline filtering: proportion of items kept/filtered and number of sends skipped due to filtering.