feat: Add validation loss tracking, early stopping, and checkpoint cleanup #2633
+315
−12
Conversation
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
This PR adds the ability to track validation loss during training: Features: - validation_fraction config option to split episodes into train/val sets - Validation loss computed using inference (select_action) for model-agnostic metrics - L1 and L2 loss metrics logged to wandb under val/ prefix - Early stopping based on validation loss or eval success rate - keep_last_n_checkpoints option to automatically cleanup old checkpoints The validation uses a separate dataset copy without augmentations for clean evaluation. Uses select_action for inference-based validation, making it policy-agnostic. Backward compatible - defaults maintain existing behavior (no validation split). Config options: - validation_fraction: 0.0-1.0 (default 0.0, no validation) - early_stopping.enable: bool (default False) - early_stopping.patience_steps: int (default 10000) - early_stopping.monitor: 'val_loss' or 'eval_success' - keep_last_n_checkpoints: int (default 0, keep all)
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
1 participant
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.
Summary
Adds optional validation loss tracking during training using a separate validation episode split.
Features
validation_fractionconfig option to split episodes into train/val setsselect_actioninference for model-agnostic metrics (L1/L2)keep_last_n_checkpointsto automatically remove old checkpointsDesign Decisions
select_actionfor validation rather than modifying individual policies - this makes validation policy-agnosticConfig Options
validation_fraction: float = 0.0 # 0.1 = 10% for validation
early_stopping.enable: bool = False
early_stopping.patience_steps: int = 10000
early_stopping.monitor: str = "val_loss" # or "eval_success"
keep_last_n_checkpoints: int = 0 # 0 = keep all
Testing