| |
| |
| |
| |
| |
| |
| |
| |
| |
| |
| |
| |
| |
| """ |
| PPO config |
| """ |
|
|
| import os |
| from dataclasses import asdict, dataclass, field, fields, is_dataclass |
| from typing import Optional, Tuple |
|
|
| from ..utils.multimodal_contract import normalize_video_source_mode |
| from ..utils.py_functional import get_abs_path |
| from ..workers.config import WorkerConfig |
|
|
|
|
| def recursive_post_init(dataclass_obj): |
| if hasattr(dataclass_obj, "post_init"): |
| dataclass_obj.post_init() |
|
|
| for attr in fields(dataclass_obj): |
| if is_dataclass(getattr(dataclass_obj, attr.name)): |
| recursive_post_init(getattr(dataclass_obj, attr.name)) |
|
|
|
|
| @dataclass |
| class DataConfig: |
| train_files: str = "" |
| val_files: str = "" |
| prompt_key: str = "prompt" |
| answer_key: str = "answer" |
| image_key: str = "images" |
| video_key: str = "videos" |
| image_dir: Optional[str] = None |
| video_fps: float = 2.0 |
| video_max_frames: int = 128 |
| max_prompt_length: int = 512 |
| max_response_length: int = 512 |
| rollout_batch_size: int = 512 |
| mini_rollout_batch_size: Optional[int] = None |
| val_batch_size: int = -1 |
| format_prompt: Optional[str] = None |
| override_chat_template: Optional[str] = None |
| enable_thinking: bool = False |
| """forwarded to the tokenizer/processor chat template; False preserves no-CoT runs""" |
| response_prefix: str = "" |
| """literal assistant prefix appended after the generation prompt, e.g. '<think>\\n'""" |
| shuffle: bool = True |
| seed: int = 1 |
| min_pixels: Optional[int] = 262144 |
| max_pixels: Optional[int] = 4194304 |
| image_min_pixels: Optional[int] = None |
| image_max_pixels: Optional[int] = None |
| video_min_pixels: Optional[int] = None |
| video_max_pixels: Optional[int] = None |
| video_total_pixels: Optional[int] = None |
| val_video_fps: Optional[float] = None |
| val_video_max_frames: Optional[int] = None |
| val_video_min_pixels: Optional[int] = None |
| val_video_max_pixels: Optional[int] = None |
| val_video_total_pixels: Optional[int] = None |
| filter_overlong_prompts: bool = True |
| filter_overlong_prompts_workers: int = 16 |
| use_preprocessed_videos: bool = True |
| """deprecated compatibility flag; prefer video_source_mode""" |
| video_source_mode: Optional[str] = None |
| """video source policy: prefer_preprocessed, preprocessed_only, realtime_only""" |
| preprocessed_video_dir: Optional[str] = None |
| """directory containing training preprocessed video files (.pt)""" |
| val_preprocessed_video_dir: Optional[str] = None |
| """directory containing validation preprocessed video files (.pt); defaults to preprocessed_video_dir""" |
| val_video_source_mode: Optional[str] = None |
| """validation video source policy; defaults to video_source_mode""" |
| inline_video_tensors: bool = False |
| """if True, decode the video once in the dataset and pass the decoded frames/metadata |
| inline through multi_modal_data so vLLM rollout and FSDP worker forward passes do |
| not re-decode the same mp4 (or reload the same .pt). Saves CPU at the cost of |
| larger pickled batches; recommended for realtime-decode runs.""" |
| group_by_task: bool = False |
| """if True, each batch contains samples from a single task type only |
| (determined by the `group_by_task_key` field in the JSONL). Avoids mixing |
| modalities (image vs video) within a batch and concentrates gradients.""" |
| group_by_task_key: str = "problem_type" |
| """JSONL field name used to identify the task type for task-grouped batching.""" |
| dataloader_num_workers: int = 8 |
| """number of subprocesses used by the training/validation StatefulDataLoader. |
| Controls how many videos are decoded in parallel per training step.""" |
|
|
| def post_init(self): |
| self.image_dir = get_abs_path(self.image_dir, prompt="Image directory") |
| self.format_prompt = get_abs_path(self.format_prompt, prompt="Format prompt file") |
| self.override_chat_template = get_abs_path(self.override_chat_template, prompt="Chat template file") |
| self.preprocessed_video_dir = get_abs_path(self.preprocessed_video_dir, prompt="Preprocessed video directory") |
| self.val_preprocessed_video_dir = get_abs_path(self.val_preprocessed_video_dir, prompt="Validation preprocessed video directory") |
| if self.image_min_pixels is None: |
| self.image_min_pixels = self.min_pixels |
| if self.image_max_pixels is None: |
| self.image_max_pixels = self.max_pixels |
| if self.video_min_pixels is None: |
| self.video_min_pixels = self.min_pixels |
| if self.video_max_pixels is None: |
| self.video_max_pixels = self.max_pixels |
| if self.val_video_fps is None: |
| self.val_video_fps = self.video_fps |
| if self.val_video_max_frames is None: |
| self.val_video_max_frames = self.video_max_frames |
| if self.val_video_min_pixels is None: |
| self.val_video_min_pixels = self.video_min_pixels |
| if self.val_video_max_pixels is None: |
| self.val_video_max_pixels = self.video_max_pixels |
| if self.val_video_total_pixels is None: |
| self.val_video_total_pixels = self.video_total_pixels |
| if self.val_preprocessed_video_dir is None: |
| self.val_preprocessed_video_dir = self.preprocessed_video_dir |
| self.video_source_mode = normalize_video_source_mode( |
| self.video_source_mode, |
| use_preprocessed_videos=self.use_preprocessed_videos, |
| ) |
| if self.val_video_source_mode is None: |
| self.val_video_source_mode = self.video_source_mode |
| else: |
| self.val_video_source_mode = normalize_video_source_mode( |
| self.val_video_source_mode, |
| use_preprocessed_videos=self.use_preprocessed_videos, |
| ) |
|
|
|
|
| @dataclass |
| class AlgorithmConfig: |
| name: str = "grpo" |
| """Algorithm selector; the released recipes are ``grpo`` and ``orarl``.""" |
| gamma: float = 1.0 |
| """discount factor for ppo gae advantage estimator""" |
| lam: float = 1.0 |
| """lambda value for ppo gae advantage estimator""" |
| adv_estimator: str = "grpo" |
| """advantage estimator, support `gae`, `grpo`, `reinforce_plus_plus`, `remax`, `rloo`""" |
| scale_rewards: bool = True |
| """Whether GRPO divides each group-centered outcome reward by the group's |
| reward standard deviation. True preserves the historical normalized GRPO |
| objective. False uses the TRL raw-centered objective |
| `A_i = r_i - mean_group(r)`, avoiding amplification of tiny reward gaps.""" |
| disable_kl: bool = False |
| """disable reference model""" |
| use_kl_loss: bool = False |
| """use kl loss instead of kl in reward""" |
| kl_penalty: str = "kl" |
| """kl penalty type, support `kl`, `abs`, `mse`, `low_var_kl`, `full`""" |
| kl_coef: float = 1e-3 |
| """kl coefficient""" |
| kl_type: str = "fixed" |
| """kl controller type, support `fixed`, `adaptive`""" |
| kl_horizon: float = 10000.0 |
| """kl horizon for adaptive kl controller""" |
| kl_target: float = 0.1 |
| """target kl for adaptive kl controller""" |
|
|
| |
| oracle_injection: bool = False |
| """Append one annotation-derived oracle response to each rollout group.""" |
| oracle_injection_mode: str = "append" |
| """Oracle insertion mode. The released OraRL recipe requires ``append``.""" |
| oracle_builder: Optional[str] = None |
| """Dotted ``module:function`` that constructs an oracle response.""" |
| oracle_replace_index: int = -1 |
| """Compatibility slot for non-released replacement mode.""" |
| oracle_log_exclude: bool = True |
| """Exclude oracle rows from on-policy reward/accuracy logging.""" |
| directional_gain: bool = False |
| """Enable the oracle-gap directional gain on on-policy utilities.""" |
| directional_gain_gamma: float = 0.25 |
| """Exponent for the clipped oracle-gap scale.""" |
| directional_gain_positive_only: bool = True |
| """Amplify only utilities pointing toward the positive oracle.""" |
| directional_gain_recenter: bool = True |
| """Re-center transformed on-policy utilities before selection.""" |
| detached_oracle_advantage: bool = False |
| """Overwrite the oracle row with a detached positive anchor.""" |
| detached_oracle_advantage_scale: float = 2.0 |
| """Base magnitude of the detached oracle anchor.""" |
| detached_oracle_use_directional_gain: bool = False |
| """Whether the policy directional gain also scales the oracle anchor.""" |
| detached_oracle_match_best_ratio: float = 1.2 |
| """Cap the oracle anchor relative to the strongest positive policy row.""" |
| detached_oracle_match_best_min: float = 0.05 |
| """Lower bound for the adaptive oracle cap.""" |
| detached_oracle_match_best_max: float = 1.0 |
| """Upper bound for the adaptive oracle cap.""" |
| oracle_reward_gate_beta: float = 2.0 |
| """Exponent for the normalized oracle reward-gap gate.""" |
| selection_prune_ratio: float = 0.0 |
| """Fraction of generated policy rows removed before actor backpropagation.""" |
| selection_keep_oracle: bool = True |
| """Force-retain the detached oracle row.""" |
| selection_positive_quota: int = 0 |
| """Number of positive-advantage policy rows retained per group.""" |
| selection_negative_quota: int = 0 |
| """Number of negative-advantage policy rows retained per group.""" |
| selection_strict_sign_balance: bool = True |
| """Select from true positive/negative buckets with deterministic fallback.""" |
| post_selection_recenter: bool = False |
| """Restore zero mean over rows that actually receive gradients.""" |
| post_selection_rms_match: bool = False |
| """Only downscale active RMS toward its pre-selection policy reference.""" |
| post_selection_rms_min_scale: float = 0.25 |
| """Lower bound for post-selection RMS downscaling.""" |
|
|
|
|
| @dataclass |
| class TrainerConfig: |
| total_epochs: int = 15 |
| """total epochs for training""" |
| max_steps: Optional[int] = None |
| """max steps for training, if specified, total_epochs is ignored""" |
| project_name: str = "orarl" |
| """project name for logger""" |
| experiment_name: str = "demo" |
| """experiment name for logger""" |
| logger: Tuple[str] = ("console", "wandb") |
| """logger type, support `console`, `mlflow`, `swanlab`, `tensorboard`, `wandb`""" |
| nnodes: int = 1 |
| """number of nodes for training""" |
| n_gpus_per_node: int = 8 |
| """number of gpus per node for training""" |
| critic_warmup: int = 0 |
| """critic warmup steps""" |
| val_freq: int = -1 |
| """validation frequency, -1 means no validation""" |
| val_before_train: bool = True |
| """validate before training""" |
| val_only: bool = False |
| """validate only, skip training""" |
| val_generations_to_log: int = 0 |
| """number of generations to log for validation""" |
| save_freq: int = -1 |
| """save frequency, -1 means no saving""" |
| save_limit: int = -1 |
| """max number of checkpoints to save, -1 means no limit""" |
| save_model_only: bool = False |
| """save model only, no optimizer state dict""" |
| keep_optim_only_latest: bool = False |
| """when True, only the latest checkpoint keeps optimizer / extra_state / dataloader; |
| older checkpoints retained within `save_limit` are thinned down to model weights only. |
| Mutually exclusive with `save_model_only` (when `save_model_only=True` optimizer is |
| never saved at all, so this flag is a no-op). Useful for saving disk while still |
| allowing resume from the latest step.""" |
| save_checkpoint_path: Optional[str] = None |
| """save checkpoint path, if not specified, use `checkpoints/project_name/experiment_name`""" |
| load_checkpoint_path: Optional[str] = None |
| """load checkpoint path""" |
| ray_timeline: Optional[str] = None |
| """file to save ray timeline""" |
| find_last_checkpoint: bool = True |
| """automatically find the last checkpoint in the save checkpoint path to resume training""" |
| keep_best_train_ckpt: bool = False |
| """When True, additionally save the model-only checkpoint at the step with the |
| highest *smoothed* training reward into a separate `best_train/` subdirectory. |
| This sidesteps the rolling `save_limit` window so a peak you discovered mid-run |
| is not evicted by later (potentially overtrained) saves. No optimizer/dataloader |
| state is kept, only model weights — purely for downstream evaluation.""" |
| best_train_metric_key: str = "reward/overall" |
| """Which key in the per-step `metrics` dict to track. Common choices: |
| `reward/overall` (default; combines IoU + format), `reward/iou`, |
| `critic/rewards/mean`.""" |
| best_train_smooth_window: int = 5 |
| """Number of recent steps to average for the best-train signal. Reduces single- |
| step noise so we lock onto a true plateau rather than a lucky spike. Set to 1 |
| to disable smoothing.""" |
| best_train_min_step: int = 10 |
| """Don't track best-train until this many steps have elapsed (skip warmup |
| where reward jumps quickly and the highest single value is nearly meaningless |
| for picking a good policy).""" |
|
|
| def post_init(self): |
| if self.save_checkpoint_path is None: |
| self.save_checkpoint_path = os.path.join("checkpoints", self.project_name, self.experiment_name) |
|
|
| self.save_checkpoint_path = os.path.abspath(self.save_checkpoint_path) |
| self.load_checkpoint_path = get_abs_path(self.load_checkpoint_path, prompt="Model checkpoint") |
|
|
|
|
| @dataclass |
| class PPOConfig: |
| data: DataConfig = field(default_factory=DataConfig) |
| worker: WorkerConfig = field(default_factory=WorkerConfig) |
| algorithm: AlgorithmConfig = field(default_factory=AlgorithmConfig) |
| trainer: TrainerConfig = field(default_factory=TrainerConfig) |
|
|
| def post_init(self): |
| self.worker.rollout.prompt_length = self.data.max_prompt_length |
| self.worker.rollout.response_length = self.data.max_response_length |
| self.worker.rollout.trust_remote_code = self.worker.actor.model.trust_remote_code |
| self.worker.actor.disable_kl = self.algorithm.disable_kl |
| self.worker.actor.use_kl_loss = self.algorithm.use_kl_loss |
| self.worker.actor.kl_penalty = self.algorithm.kl_penalty |
| self.worker.actor.kl_coef = self.algorithm.kl_coef |
| |
| |
| self.worker.actor.selection_prune_ratio = ( |
| float(self.algorithm.selection_prune_ratio) |
| if str(self.algorithm.name).strip().lower() == "orarl" |
| else 0.0 |
| ) |
|
|
| def deep_post_init(self): |
| recursive_post_init(self) |
|
|
| def to_dict(self): |
| from .orarl_config import public_config_dict |
|
|
| return public_config_dict(self) |
|
|
