From f5d54d5e5dc63b4b3b3279efc48e2c78657145aa Mon Sep 17 00:00:00 2001 From: Varad Pimpalkhute <35189598+nightlessbaron@users.noreply.github.com> Date: Wed, 16 Sep 2026 14:31:54 -0700 Subject: [PATCH] feat: add QAE and adaptive PPO clipping Add quantile reward baselines for GRPO/GSPO and the MAI-Thinking-1 entropy-feedback controller for the PPO upper clipping bound. Persist controller state across checkpoints and support Megatron and FSDP training. --- docs/en/examples/qae-adaptive-clip.md | 46 ++++++++ .../backends/experimental/fsdp_utils/actor.py | 8 +- .../experimental/fsdp_utils/checkpoint.py | 3 + miles/backends/megatron_utils/model.py | 15 ++- miles/backends/training_utils/loss.py | 75 +++++++++++++ miles/ray/rollout.py | 13 +-- miles/utils/arguments.py | 94 +++++++++++++++- miles/utils/entropy_control.py | 101 ++++++++++++++++++ miles/utils/ppo_utils.py | 45 ++++++++ tests/fast/utils/test_entropy_control.py | 56 ++++++++++ tests/fast/utils/test_ppo_utils.py | 46 +++++++- 11 files changed, 489 insertions(+), 13 deletions(-) create mode 100644 docs/en/examples/qae-adaptive-clip.md create mode 100644 miles/utils/entropy_control.py create mode 100644 tests/fast/utils/test_entropy_control.py diff --git a/docs/en/examples/qae-adaptive-clip.md b/docs/en/examples/qae-adaptive-clip.md new file mode 100644 index 00000000000..7de7d26481c --- /dev/null +++ b/docs/en/examples/qae-adaptive-clip.md @@ -0,0 +1,46 @@ +# QAE with adaptive clipping + +This recipe combines Quantile Advantage Estimation (QAE) with the adaptive +upper clipping controller from MAI-Thinking-1. + +```bash +python train.py \ + ... \ + --advantage-estimator grpo \ + --qae-quantile 0.4 \ + --eps-clip 0.6 \ + --use-adaptive-clip \ + --adaptive-clip-target-entropy 0.3 \ + --adaptive-clip-step-size 0.25 \ + --adaptive-clip-max-relaxation 2.5 \ + --entropy-coef 0 +``` + +`--qae-quantile 0.4` changes only the group reward baseline. The implementation +uses the paper's right-continuous empirical quantile. For binary rewards this +produces the two QAE regimes exactly: hard groups update successes and easy +groups update failures. + +With adaptive clipping, the PPO ratio interval at step `t` is + +```text +[1 - eps_clip, 1 / (1 - eps_clip) + k_t] +``` + +and the controller applies + +```text +k_{t+1} = clip(k_t + delta * sign(target_entropy - estimated_entropy), 0, k_max). +``` + +The entropy estimate is the importance-weighted all-token average from the MAI +report. The controller updates once after each optimizer step and its state is +saved beside checkpoints. Monitor these training metrics: + +- `adaptive_clip_entropy`: controller entropy estimate +- `adaptive_clip_relaxation`: current `k` +- `adaptive_clip_high`: MILES upper clip offset (the upper ratio is `1 + value`) + +The MAI values above are a starting point, not universal defaults for every +model and rollout temperature. A sign controller commonly oscillates in a +narrow band around the target; it does not guarantee a smooth sine wave. diff --git a/miles/backends/experimental/fsdp_utils/actor.py b/miles/backends/experimental/fsdp_utils/actor.py index 47c2540f98d..6c12af36c36 100644 --- a/miles/backends/experimental/fsdp_utils/actor.py +++ b/miles/backends/experimental/fsdp_utils/actor.py @@ -29,7 +29,12 @@ log_rollout_data, log_train_step, ) -from ...training_utils.loss import compute_advantages_and_returns, get_log_probs_and_entropy, loss_function +from ...training_utils.loss import ( + compute_advantages_and_returns, + get_log_probs_and_entropy, + loss_function, + update_adaptive_clip, +) from ...training_utils.parallel import get_parallel_state, set_parallel_state from . import checkpoint from .lr_scheduler import get_lr_scheduler @@ -492,6 +497,7 @@ def _train_core(self, rollout_id: int, rollout_data) -> None: ) loss_dict = aggregate_train_losses(losses_reduced) + update_adaptive_clip(self.args, loss_dict) extra_metrics = {} for param_group_id, param_group in enumerate(self.optimizer.param_groups): diff --git a/miles/backends/experimental/fsdp_utils/checkpoint.py b/miles/backends/experimental/fsdp_utils/checkpoint.py index 80eeb64110a..82cede626a9 100644 --- a/miles/backends/experimental/fsdp_utils/checkpoint.py +++ b/miles/backends/experimental/fsdp_utils/checkpoint.py @@ -13,6 +13,7 @@ from torch.distributed.checkpoint.stateful import Stateful from miles.backends.training_utils.log_utils import init_train_step_counter, save_train_step_counter +from miles.utils.entropy_control import load_adaptive_clip_state, save_adaptive_clip_state logger = logging.getLogger(__name__) @@ -185,6 +186,7 @@ def finalize_load(actor: Any, checkpoint_payload: dict[str, Any] | None) -> None if getattr(actor.args, "start_rollout_id", None) is None: actor.args.start_rollout_id = iteration + load_adaptive_clip_state(actor.args, actor.args.load, iteration) if dist.get_rank() == 0: init_train_step_counter(actor.args.load, iteration) @@ -251,6 +253,7 @@ def save(actor: Any, iteration: int) -> None: tracker_file = base_dir / "latest_checkpointed_iteration.txt" tracker_file.write_text(str(step_id)) save_train_step_counter(actor.args.save, step_id) + save_adaptive_clip_state(actor.args, actor.args.save, step_id) logger.info(f"[FSDP] Saved checkpoint to {checkpoint_dir}") dist.barrier() diff --git a/miles/backends/megatron_utils/model.py b/miles/backends/megatron_utils/model.py index 316a1d2975c..4cf8f24ad3e 100644 --- a/miles/backends/megatron_utils/model.py +++ b/miles/backends/megatron_utils/model.py @@ -22,6 +22,7 @@ from megatron.training.training import get_model, preprocess_common_state_dict from miles.utils.dumper_utils import DumperMegatronUtil, DumperPhase +from miles.utils.entropy_control import load_adaptive_clip_state, save_adaptive_clip_state from miles.utils.memory_utils import clear_memory from ..training_utils.ci_utils import check_grad_norm, check_kl @@ -33,7 +34,7 @@ log_train_step, save_train_step_counter, ) -from ..training_utils.loss import loss_function +from ..training_utils.loss import loss_function, update_adaptive_clip from ..training_utils.parallel import get_parallel_state from .checkpoint import load_checkpoint, save_checkpoint, save_checkpoint_with_lora from .ci_utils import ( @@ -626,6 +627,15 @@ def train( num_microbatches[step_id], ) + role = getattr(model[0], "role", "actor") + if role == "actor": + update_adaptive_clip( + args, + loss_dict, + pipeline_group=mpu.get_pipeline_model_parallel_group(), + pipeline_source_rank=mpu.get_pipeline_model_parallel_last_rank(), + ) + if step_id == 0: # Enable forward pre-hook after training step has successfully run. All subsequent # forward passes will use the forward pre-hook / `param_sync_func` in @@ -659,7 +669,6 @@ def train( # per train step log. if is_megatron_main_rank(): accumulated_step_id = rollout_id * num_steps_per_rollout + step_id - role = getattr(model[0], "role", "actor") role_tag = "" if role == "actor" else f"{role}-" extra_metrics = {} @@ -741,6 +750,7 @@ def save( clear_memory() if is_megatron_main_rank(): save_train_step_counter(args.save, iteration) + save_adaptive_clip_state(args, args.save, iteration) if hashes is not None: save_model_hashes(args, model, iteration, hashes) @@ -840,6 +850,7 @@ def initialize_model_and_optimizer( clear_memory() check_model_hashes(args, model, iteration) + load_adaptive_clip_state(args, args.load, iteration) # A full resume restores the scheduler in load_checkpoint(). Advancing it # again here would double-count the completed iterations. Only reconstruct diff --git a/miles/backends/training_utils/loss.py b/miles/backends/training_utils/loss.py index 625bcd33e6c..634daabb7ea 100644 --- a/miles/backends/training_utils/loss.py +++ b/miles/backends/training_utils/loss.py @@ -8,11 +8,13 @@ from torch.utils.checkpoint import checkpoint from miles.utils.distributed_utils import distributed_masked_whiten +from miles.utils.entropy_control import adaptive_clip_high, update_adaptive_clip_relaxation from miles.utils.misc import load_function from miles.utils.ppo_utils import ( calculate_log_probs_and_entropy, compute_approx_kl, compute_gspo_kl, + compute_importance_weighted_entropy, compute_opd_reward, compute_opsm_mask, compute_policy_loss, @@ -34,6 +36,58 @@ logger = logging.getLogger(__name__) +def update_adaptive_clip( + args: Namespace, + loss_dict: dict[str, float], + *, + pipeline_group: torch.distributed.ProcessGroup | None = None, + pipeline_source_rank: int | None = None, +) -> None: + """Update and synchronize MAI's clipping controller after an optimizer step. + + The policy loss emits an entropy numerator and token count for each + microbatch. ``aggregate_train_losses`` sums both globally; their ratio is + therefore the paper's all-token importance-weighted entropy estimate. + """ + if not getattr(args, "use_adaptive_clip", False): + return + + parallel_state = get_parallel_state() + is_last_stage = parallel_state.is_pp_last_stage + state = torch.tensor( + [float(args.adaptive_clip_relaxation), float("nan")], + dtype=torch.float64, + device=torch.cuda.current_device(), + ) + if is_last_stage: + entropy_sum = loss_dict.pop("adaptive_clip_entropy_sum") + token_count = loss_dict.pop("adaptive_clip_token_count") + if token_count <= 0: + raise RuntimeError("adaptive clipping received a training step with no valid response tokens") + estimated_entropy = entropy_sum / token_count + relaxation = update_adaptive_clip_relaxation( + args.adaptive_clip_relaxation, + estimated_entropy, + args.adaptive_clip_target_entropy, + args.adaptive_clip_step_size, + args.adaptive_clip_max_relaxation, + ) + state[0] = relaxation + state[1] = estimated_entropy + + if pipeline_group is not None and torch.distributed.get_world_size(pipeline_group) > 1: + if pipeline_source_rank is None: + raise ValueError("pipeline_source_rank is required when synchronizing adaptive clipping") + torch.distributed.broadcast(state, src=pipeline_source_rank, group=pipeline_group) + + args.adaptive_clip_relaxation = state[0].item() + args.eps_clip_high = adaptive_clip_high(args.eps_clip, args.adaptive_clip_relaxation) + if is_last_stage: + loss_dict["adaptive_clip_entropy"] = state[1].item() + loss_dict["adaptive_clip_relaxation"] = args.adaptive_clip_relaxation + loss_dict["adaptive_clip_high"] = args.eps_clip_high + + def _nan_dbg_flush_logs() -> None: """Flush logger/stdout/stderr so crash diagnostics survive abrupt failures.""" for handler in logging.getLogger().handlers + logger.handlers: @@ -763,6 +817,26 @@ def policy_loss_function( extra=f"ratio_delta_min={ratio_delta_min} ratio_delta_max={ratio_delta_max}", ) + adaptive_clip_metrics = {} + if getattr(args, "use_adaptive_clip", False): + # MAI Eq. (7): mean over every valid response token of + # -log pi_theta(y_t | x, y_ list[EvalDatasetConfig]: return eval_datasets +def _validate_entropy_control_args(args) -> None: + if args.qae_quantile is not None: + if not 0 < args.qae_quantile < 1: + raise ValueError(f"--qae-quantile must be in (0, 1), got {args.qae_quantile}") + if args.advantage_estimator not in ["grpo", "gspo"]: + raise ValueError("--qae-quantile requires --advantage-estimator grpo or gspo") + if not args.rewards_normalization: + raise ValueError("--qae-quantile requires reward normalization") + if args.n_samples_per_prompt < 2: + raise ValueError("--qae-quantile requires at least two samples per prompt") + if ( + args.custom_reward_post_process_path is not None + or args.custom_convert_samples_to_train_data_path is not None + ): + raise ValueError( + "--qae-quantile uses MILES' built-in reward post-processing and is incompatible with custom " + "reward post-processing or sample conversion" + ) + + if args.use_adaptive_clip: + if args.advantage_estimator not in ["grpo", "gspo"]: + raise ValueError("--use-adaptive-clip requires --advantage-estimator grpo or gspo") + if args.loss_type != "policy_loss": + raise ValueError("--use-adaptive-clip requires --loss-type policy_loss") + if args.eps_clip_high is not None: + raise ValueError("--eps-clip-high is controlled by --use-adaptive-clip and must be unset") + if not 0 <= args.eps_clip < 1: + raise ValueError(f"adaptive clipping requires --eps-clip in [0, 1), got {args.eps_clip}") + if args.adaptive_clip_target_entropy < 0: + raise ValueError("--adaptive-clip-target-entropy must be non-negative") + if args.adaptive_clip_step_size <= 0: + raise ValueError("--adaptive-clip-step-size must be positive") + if args.adaptive_clip_max_relaxation < 0: + raise ValueError("--adaptive-clip-max-relaxation must be non-negative") + if not 0 <= args.adaptive_clip_initial_relaxation <= args.adaptive_clip_max_relaxation: + raise ValueError( + "--adaptive-clip-initial-relaxation must be between zero and " + "--adaptive-clip-max-relaxation" + ) + if args.entropy_coef != 0: + logger.warning( + "Adaptive clipping was evaluated without an explicit entropy bonus; " + "consider setting --entropy-coef 0." + ) + initialize_adaptive_clip(args) + elif args.eps_clip_high is None: + args.eps_clip_high = args.eps_clip + + def miles_validate_args(args): args.eval_datasets = _resolve_eval_datasets(args) @@ -1986,9 +2075,6 @@ def miles_validate_args(args): if getattr(args, "balance_by_flops", False): assert args.use_dynamic_batch_size, "--balance-by-flops requires --use-dynamic-batch-size" - if args.eps_clip_high is None: - args.eps_clip_high = args.eps_clip - if args.eval_reward_key is None: args.eval_reward_key = args.reward_key @@ -2120,6 +2206,8 @@ def miles_validate_args(args): logger.info(f"Warning: Argument {k} is already set to {getattr(args, k)}, will override with {v}.") setattr(args, k, v) + _validate_entropy_control_args(args) + if args.eval_max_context_len is None: logger.info( f"args.eval_max_context_len is not set. Use args.rollout_max_context_len {args.rollout_max_context_len} as default value." diff --git a/miles/utils/entropy_control.py b/miles/utils/entropy_control.py new file mode 100644 index 00000000000..a742d35f93b --- /dev/null +++ b/miles/utils/entropy_control.py @@ -0,0 +1,101 @@ +"""Entropy-control helpers shared by the RL training backends. + +The adaptive clipping rule follows MAI-Thinking-1, Section 3.1.1. MILES +stores clip ranges as offsets around one, so the paper's upper ratio bound + + (1 - epsilon) ** -1 + k + +is represented by ``eps_clip_high = (1 - epsilon) ** -1 - 1 + k``. +""" + +from __future__ import annotations + +import json +import logging +from argparse import Namespace +from pathlib import Path + + +logger = logging.getLogger(__name__) + +_STATE_FILE = "adaptive_clip_state.json" + + +def adaptive_clip_high(eps_clip: float, relaxation: float) -> float: + """Return MILES' upper clip offset for MAI's log-symmetric interval.""" + if not 0 <= eps_clip < 1: + raise ValueError(f"eps_clip must be in [0, 1), got {eps_clip}") + if relaxation < 0: + raise ValueError(f"adaptive clip relaxation must be non-negative, got {relaxation}") + return 1.0 / (1.0 - eps_clip) - 1.0 + relaxation + + +def update_adaptive_clip_relaxation( + relaxation: float, + estimated_entropy: float, + target_entropy: float, + step_size: float, + max_relaxation: float, +) -> float: + """Apply MAI's sign-integral controller update to ``k``.""" + if estimated_entropy < target_entropy: + direction = 1.0 + elif estimated_entropy > target_entropy: + direction = -1.0 + else: + direction = 0.0 + return min(max(relaxation + step_size * direction, 0.0), max_relaxation) + + +def initialize_adaptive_clip(args: Namespace) -> None: + """Initialize the runtime controller state from command-line arguments.""" + if not getattr(args, "use_adaptive_clip", False): + return + args.adaptive_clip_relaxation = float(args.adaptive_clip_initial_relaxation) + args.eps_clip_high = adaptive_clip_high(args.eps_clip, args.adaptive_clip_relaxation) + + +def _checkpoint_state_path(checkpoint_dir: str | None, iteration: int | None) -> Path | None: + if checkpoint_dir is None or iteration is None: + return None + base = Path(checkpoint_dir).expanduser() + iteration_dir = f"iter_{int(iteration):07d}" + return (base if base.name == iteration_dir else base / iteration_dir) / _STATE_FILE + + +def save_adaptive_clip_state(args: Namespace, checkpoint_dir: str | None, iteration: int | None) -> None: + """Persist controller state next to a model checkpoint.""" + if not getattr(args, "use_adaptive_clip", False): + return + path = _checkpoint_state_path(checkpoint_dir, iteration) + if path is None: + return + payload = { + "relaxation": float(args.adaptive_clip_relaxation), + "eps_clip_high": float(args.eps_clip_high), + } + try: + path.parent.mkdir(parents=True, exist_ok=True) + temporary_path = path.with_suffix(path.suffix + ".tmp") + temporary_path.write_text(json.dumps(payload, sort_keys=True) + "\n") + temporary_path.replace(path) + except OSError as error: + logger.warning("Failed to persist adaptive-clip state: %s", error) + + +def load_adaptive_clip_state(args: Namespace, checkpoint_dir: str | None, iteration: int | None) -> None: + """Restore controller state, retaining the configured initial state if absent.""" + if not getattr(args, "use_adaptive_clip", False): + return + path = _checkpoint_state_path(checkpoint_dir, iteration) + if path is None or not path.is_file(): + return + try: + payload = json.loads(path.read_text()) + relaxation = float(payload["relaxation"]) + except (OSError, KeyError, TypeError, ValueError, json.JSONDecodeError) as error: + logger.warning("Failed to restore adaptive-clip state from %s: %s", path, error) + return + + args.adaptive_clip_relaxation = min(max(relaxation, 0.0), args.adaptive_clip_max_relaxation) + args.eps_clip_high = adaptive_clip_high(args.eps_clip, args.adaptive_clip_relaxation) diff --git a/miles/utils/ppo_utils.py b/miles/utils/ppo_utils.py index 9e04c142aac..8e9fa3f8cbf 100644 --- a/miles/utils/ppo_utils.py +++ b/miles/utils/ppo_utils.py @@ -1,6 +1,7 @@ # Adapt from https://github.com/OpenRLHF/OpenRLHF/blob/10c733694ed9fbb78a0a2ff6a05efc7401584d46/openrlhf/models/utils.py # and https://github.com/OpenRLHF/OpenRLHF/blob/10c733694ed9fbb78a0a2ff6a05efc7401584d46/openrlhf/trainer/ppo_utils/experience_maker.py +import math from argparse import Namespace import torch @@ -8,6 +9,50 @@ import torch.nn.functional as F +def compute_group_advantages( + rewards: torch.Tensor, + *, + quantile: float | None = None, + normalize_by_std: bool = True, + epsilon: float = 1e-6, +) -> tuple[torch.Tensor, torch.Tensor]: + """Center grouped rewards with either the GRPO mean or the QAE baseline. + + ``rewards`` must have shape ``[num_prompts, samples_per_prompt]``. When + ``quantile`` is set, the baseline is the right-continuous empirical + quantile from QAE, i.e. order statistic ``ceil(K * G)``. This is not + ``torch.quantile``'s default linear interpolation; for binary rewards the + order statistic is what gives QAE's exact zero/one, two-regime gate. + """ + if rewards.ndim != 2: + raise ValueError(f"grouped rewards must be rank 2, got shape {tuple(rewards.shape)}") + if rewards.size(1) == 0: + raise ValueError("each reward group must contain at least one sample") + + if quantile is None: + baseline = rewards.mean(dim=-1, keepdim=True) + else: + if not 0 < quantile < 1: + raise ValueError(f"QAE quantile must be in (0, 1), got {quantile}") + order = math.ceil(quantile * rewards.size(1)) + baseline = rewards.kthvalue(order, dim=-1, keepdim=True).values + + advantages = rewards - baseline + if normalize_by_std: + advantages = advantages / (rewards.std(dim=-1, keepdim=True) + epsilon) + return advantages, baseline + + +def compute_importance_weighted_entropy( + log_probs: torch.Tensor, + old_log_probs: torch.Tensor, +) -> torch.Tensor: + """Return the per-token entropy estimator from MAI-Thinking-1 Eq. (7).""" + log_probs = log_probs.detach().float() + old_log_probs = old_log_probs.detach().float() + return -log_probs * torch.exp(log_probs - old_log_probs) + + @torch.compile(dynamic=True) def compute_approx_kl( log_probs: torch.Tensor, diff --git a/tests/fast/utils/test_entropy_control.py b/tests/fast/utils/test_entropy_control.py new file mode 100644 index 00000000000..7468741362d --- /dev/null +++ b/tests/fast/utils/test_entropy_control.py @@ -0,0 +1,56 @@ +from argparse import Namespace + +import pytest + +from miles.utils.entropy_control import ( + adaptive_clip_high, + initialize_adaptive_clip, + load_adaptive_clip_state, + save_adaptive_clip_state, + update_adaptive_clip_relaxation, +) + + +def test_adaptive_clip_high_matches_mai_parameterization() -> None: + assert adaptive_clip_high(eps_clip=0.6, relaxation=0.0) == pytest.approx(1.5) + assert adaptive_clip_high(eps_clip=0.6, relaxation=0.25) == pytest.approx(1.75) + + +@pytest.mark.parametrize( + ("entropy", "expected"), + [(0.2, 0.75), (0.3, 0.5), (0.4, 0.25)], +) +def test_adaptive_clip_controller_moves_toward_target(entropy: float, expected: float) -> None: + actual = update_adaptive_clip_relaxation( + relaxation=0.5, + estimated_entropy=entropy, + target_entropy=0.3, + step_size=0.25, + max_relaxation=1.0, + ) + assert actual == pytest.approx(expected) + + +def test_adaptive_clip_controller_clamps_relaxation() -> None: + assert update_adaptive_clip_relaxation(0.0, 1.0, 0.3, 0.25, 1.0) == 0.0 + assert update_adaptive_clip_relaxation(1.0, 0.0, 0.3, 0.25, 1.0) == 1.0 + + +def test_adaptive_clip_state_round_trip(tmp_path) -> None: + args = Namespace( + use_adaptive_clip=True, + eps_clip=0.6, + adaptive_clip_initial_relaxation=0.0, + adaptive_clip_max_relaxation=2.5, + ) + initialize_adaptive_clip(args) + args.adaptive_clip_relaxation = 0.75 + args.eps_clip_high = adaptive_clip_high(args.eps_clip, args.adaptive_clip_relaxation) + + save_adaptive_clip_state(args, str(tmp_path), 3) + args.adaptive_clip_relaxation = 0.0 + args.eps_clip_high = 0.0 + load_adaptive_clip_state(args, str(tmp_path), 3) + + assert args.adaptive_clip_relaxation == pytest.approx(0.75) + assert args.eps_clip_high == pytest.approx(2.25) diff --git a/tests/fast/utils/test_ppo_utils.py b/tests/fast/utils/test_ppo_utils.py index 4b2e9ce9a11..ba2eb7fbeda 100644 --- a/tests/fast/utils/test_ppo_utils.py +++ b/tests/fast/utils/test_ppo_utils.py @@ -1,7 +1,51 @@ import pytest import torch -from miles.utils.ppo_utils import compute_opd_reward +from miles.utils.ppo_utils import compute_group_advantages, compute_importance_weighted_entropy, compute_opd_reward + + +def test_compute_group_advantages_preserves_grpo_mean_baseline() -> None: + rewards = torch.tensor([[0.0, 1.0, 3.0], [2.0, 4.0, 6.0]]) + + advantages, baseline = compute_group_advantages(rewards, normalize_by_std=False) + + torch.testing.assert_close(baseline, rewards.mean(dim=-1, keepdim=True)) + torch.testing.assert_close(advantages, rewards - rewards.mean(dim=-1, keepdim=True)) + + +def test_compute_group_advantages_qae_uses_right_continuous_order_statistic() -> None: + # At K=0.5 with four binary samples, the empirical inverse CDF selects the + # second order statistic (zero), not torch.quantile's interpolated 0.5. + rewards = torch.tensor([[0.0, 0.0, 1.0, 1.0]]) + + advantages, baseline = compute_group_advantages(rewards, quantile=0.5, normalize_by_std=False) + + torch.testing.assert_close(baseline, torch.tensor([[0.0]])) + torch.testing.assert_close(advantages, torch.tensor([[0.0, 0.0, 1.0, 1.0]])) + + +def test_compute_group_advantages_qae_selects_only_failures_for_easy_group() -> None: + rewards = torch.tensor([[0.0, 0.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0]]) + + advantages, baseline = compute_group_advantages(rewards, quantile=0.4, normalize_by_std=False) + + torch.testing.assert_close(baseline, torch.tensor([[1.0]])) + torch.testing.assert_close(advantages, rewards - 1.0) + + +def test_compute_group_advantages_rejects_invalid_quantile() -> None: + with pytest.raises(ValueError, match="QAE quantile"): + compute_group_advantages(torch.tensor([[0.0, 1.0]]), quantile=1.0) + + +def test_compute_importance_weighted_entropy() -> None: + log_probs = torch.log(torch.tensor([0.25, 0.5])) + old_log_probs = torch.log(torch.tensor([0.5, 0.25])) + expected = -log_probs * torch.tensor([0.5, 2.0]) + + actual = compute_importance_weighted_entropy(log_probs, old_log_probs) + + torch.testing.assert_close(actual, expected) def test_compute_opd_reward_logr() -> None: