ViperEkura/AstrAI
1
0
Fork 0
Code Issues Pull Requests Packages Projects Releases Wiki Activity

feat: add online ppo with value-model critic and gae advantages

Browse Source 
- register online_ppo train type backed by PPOStrategy: token-level clipped surrogate over GAE advantages plus masked value regression against rollout-pinned returns, with explained-variance metrics
- fold the reference-KL penalty (k3 estimator) into per-token rewards before GAE and pin advantages/returns on RolloutResult so replayed gradient steps optimize fixed targets
- add self-contained ValueModel critic with a zero-initialized value head and backbone warm-started from policy weights; AutoRegressiveLM stays untouched and trunk parity is pinned by tests
- step the critic's own optimizer outside the policy-version lock with the same max_grad_norm clipping as the policy
- persist critic state as value_model.pt/value_optimizer.pt checkpoint extras; resume restores it, fails loudly when missing, and the train.sh completeness check requires the extras for online_ppo configs
- extract shared rollout sequence/logprob helpers from GRPO (behavior unchanged) and add ppo_gamma/ppo_gae_lambda/ppo_vf_coef CLI options
This commit is contained in:
ViperEkura
2026-09-05 01:59:50 +08:00
parent 816b96a58a
commit 350e4a1849
17 changed files with 1390 additions and 72 deletions
Download Patch File Download Diff File Expand all files Collapse all files
tests/trainer/test_ppo_strategy.py
+469
View File
@@ -0,0 +1,469 @@
"""Unit tests for PPO: GAE numerics, the ValueModel critic, and PPOStrategy."""
import pytest
import torch
import astrai.trainer.strategy as strategy_module
from astrai.model.transformer import AutoRegressiveLM
from astrai.model.value import ValueModel
from astrai.trainer.rollout import RolloutResult
from astrai.trainer.strategy import (
PPOStrategy,
StrategyFactory,
compute_gae,
)
from tests.helpers import FakeExecutor, make_frozen, make_model, make_rollout_config
def _make_batch(
batch_size=2, group_size=4, prompt_len=8, response_len=12, device="cpu"
):
"""Construct a PPO batch with deterministic shapes.
Returns dict with prompts [B, P], responses [B, G, R], masks [B, G, R],
rewards [B, G], logprobs_old [B, G, R].
"""
return {
"prompts": torch.randint(0, 200, (batch_size, prompt_len), device=device),
"responses": torch.randint(
0, 200, (batch_size, group_size, response_len), device=device
),
"masks": torch.ones(batch_size, group_size, response_len, device=device),
"rewards": torch.randn(batch_size, group_size, device=device),
"logprobs_old": torch.zeros(
batch_size, group_size, response_len, device=device
),
}
def _make_value_model(policy_model, device):
"""Build a ValueModel whose backbone warm-starts from the policy."""
critic = ValueModel(policy_model.config).to(device=device)
result = critic.load_state_dict(policy_model.state_dict(), strict=False)
assert not result.unexpected_keys
assert all(key.startswith("value_head.") for key in result.missing_keys)
return critic
@pytest.fixture
def ppo_strategy(device):
model, _ = make_model(device)
critic = _make_value_model(model, device)
strategy = PPOStrategy(
model=model,
device=device,
critic=critic,
critic_optimizer=torch.optim.AdamW(critic.parameters(), lr=1e-3),
ref_model=make_frozen(model, device),
clip_eps=0.2,
kl_coef=0.01,
gamma=1.0,
gae_lambda=0.95,
vf_coef=0.5,
executor=FakeExecutor(),
)
return strategy, device
# ============== compute_gae ==============
def test_gae_monte_carlo_when_values_zero(device):
"""γ=1, λ=1, V=0: advantage and return equal the terminal reward at
every valid position (Monte-Carlo return)."""
B, G, R = 2, 3, 4
rewards = torch.zeros(B, G, R, device=device)
rewards[..., -1] = 1.0
values = torch.zeros(B, G, R, device=device)
mask = torch.ones(B, G, R, dtype=torch.bool, device=device)
advantages, returns = compute_gae(rewards, values, mask, gamma=1.0, gae_lambda=1.0)
assert torch.allclose(advantages, torch.full_like(rewards, 1.0))
assert torch.allclose(returns, torch.full_like(rewards, 1.0))
def test_gae_lambda_zero_is_one_step_td(device):
"""λ=0: advantage degenerates to the TD residual δ_t."""
torch.manual_seed(0)
rewards = torch.zeros(1, 1, 3, device=device)
rewards[0, 0, -1] = 2.0
values = torch.tensor([[[0.5, 1.0, -0.5]]], device=device)
mask = torch.ones(1, 1, 3, dtype=torch.bool, device=device)
advantages, returns = compute_gae(rewards, values, mask, gamma=0.9, gae_lambda=0.0)
# δ_2 = r + 0 - V_2 = 2.5; δ_1 = 0 + 0.9·V_2 - V_1 = -1.45;
# δ_0 = 0 + 0.9·V_1 - V_0 = 0.4
expected = torch.tensor([[[0.4, -1.45, 2.5]]], device=device)
assert torch.allclose(advantages, expected, atol=1e-6)
assert torch.allclose(returns, advantages + values, atol=1e-6)
def test_gae_hand_computed_discounted_case(device):
"""γ=0.9, λ=0.8 against a hand-rolled backward accumulation."""
rewards = torch.tensor([[[0.0, 0.0, 1.0]]], device=device)
values = torch.tensor([[[0.1, 0.2, 0.3]]], device=device)
mask = torch.ones(1, 1, 3, dtype=torch.bool, device=device)
gamma, lam = 0.9, 0.8
advantages, returns = compute_gae(rewards, values, mask, gamma, lam)
delta2 = 1.0 + 0.0 - 0.3
gae2 = delta2
delta1 = 0.0 + gamma * 0.3 - 0.2
gae1 = delta1 + gamma * lam * gae2
delta0 = 0.0 + gamma * 0.2 - 0.1
gae0 = delta0 + gamma * lam * gae1
expected = torch.tensor([[[gae0, gae1, gae2]]], device=device)
assert torch.allclose(advantages, expected, atol=1e-6)
assert torch.allclose(returns, expected + values, atol=1e-6)
def test_gae_padding_does_not_leak(device):
"""Garbage values at padded positions must not change valid outputs."""
torch.manual_seed(1)
B, G, R = 2, 2, 5
rewards = torch.zeros(B, G, R, device=device)
rewards[0, 0, 2] = 1.0 # terminal at position 2 of a length-3 response
values = torch.randn(B, G, R, device=device)
mask = torch.ones(B, G, R, dtype=torch.bool, device=device)
mask[0, 0, 3:] = False
mask[1, :, 2:] = False
rewards[0, 0, 3:] = 100.0 # reward garbage in padding must be ignored
advantages, returns = compute_gae(rewards, values, mask, gamma=0.9, gae_lambda=0.9)
assert torch.allclose(advantages[0, 0, 3:], torch.zeros_like(advantages[0, 0, 3:]))
assert torch.allclose(returns[0, 0, 3:], torch.zeros_like(returns[0, 0, 3:]))
# The terminal reward at position 2 still drives a finite advantage.
assert advantages[0, 0, 2] != 0.0
def test_gae_empty_response_is_all_zero(device):
"""A fully padded response yields zero advantages and returns."""
rewards = torch.zeros(1, 1, 3, device=device)
values = torch.randn(1, 1, 3, device=device)
mask = torch.zeros(1, 1, 3, dtype=torch.bool, device=device)
advantages, returns = compute_gae(rewards, values, mask, 1.0, 0.95)
assert torch.count_nonzero(advantages) == 0
assert torch.count_nonzero(returns) == 0
# ============== ValueModel ==============
def test_value_model_trunk_matches_policy_hidden_states(device):
"""ValueModel's trunk reproduces AutoRegressiveLM's hidden states.
Pins the duplicated trunk pass in ``ValueModel.forward`` to the policy
forward: a ones-initialized value head must return the row-wise sum of
the policy's ``hidden_states``.
"""
model, _ = make_model(device)
critic = _make_value_model(model, device)
with torch.no_grad():
critic.value_head.weight.fill_(1.0)
critic.value_head.bias.zero_()
torch.manual_seed(2)
input_ids = torch.randint(0, 200, (2, 10), device=device)
input_mask = torch.ones(2, 10, dtype=torch.bool, device=device)
input_mask[1, :3] = False
with torch.no_grad():
policy_hidden = model(input_ids, input_mask=input_mask)["hidden_states"]
values = critic(input_ids, input_mask=input_mask)["values"]
assert values.shape == (2, 10)
assert torch.allclose(values, policy_hidden.sum(dim=-1), atol=1e-5)
def test_value_model_zero_head_outputs_zero(device):
model, _ = make_model(device)
critic = _make_value_model(model, device)
input_ids = torch.randint(0, 200, (2, 8), device=device)
with torch.no_grad():
values = critic(input_ids)["values"]
assert torch.count_nonzero(values) == 0
def test_value_model_rejects_packed_inference_input(device):
critic = ValueModel(make_rollout_config()).to(device=device)
with pytest.raises(ValueError, match="critic input_ids"):
critic(torch.randint(0, 200, (16,), device=device))
# ============== PPOStrategy ==============
def test_factory_registers_online_ppo():
assert StrategyFactory.is_registered("online_ppo")
assert StrategyFactory.get_component_class("online_ppo") is PPOStrategy
def test_ppo_supports_online(ppo_strategy):
strategy, _ = ppo_strategy
assert strategy.supports_online() is True
def test_ppo_loss_is_finite_and_differentiable(ppo_strategy):
strategy, device = ppo_strategy
batch = _make_batch(device=device)
loss = strategy.compute_loss(batch)
assert loss.dim() == 0
assert torch.isfinite(loss).item()
loss.backward()
assert any(
p.grad is not None and p.grad.abs().sum().item() > 0
for p in strategy.model.parameters()
)
assert any(
p.grad is not None and p.grad.abs().sum().item() > 0
for p in strategy.critic.parameters()
)
def test_ppo_requires_behavior_logprobs(ppo_strategy):
strategy, device = ppo_strategy
batch = _make_batch(device=device)
del batch["logprobs_old"]
with pytest.raises(ValueError, match="logprobs_old"):
strategy.compute_loss(batch)
@pytest.mark.parametrize("invalid", ["shape", "nonfinite"])
def test_ppo_rejects_invalid_behavior_logprobs(ppo_strategy, invalid):
strategy, device = ppo_strategy
batch = _make_batch(device=device)
if invalid == "shape":
batch["logprobs_old"] = torch.zeros(1, device=device)
match = "shape must match responses"
else:
batch["logprobs_old"] = torch.zeros_like(batch["responses"], dtype=torch.float)
batch["logprobs_old"][0, 0, 0] = float("nan")
match = "only finite values"
with pytest.raises(ValueError, match=match):
strategy.compute_loss(batch)
def test_ppo_ref_model_not_updated_by_backward(ppo_strategy):
strategy, device = ppo_strategy
loss = strategy.compute_loss(_make_batch(device=device))
loss.backward()
for p in strategy.ref_model.parameters():
assert p.grad is None
def test_ppo_zero_advantage_and_zero_critic_gives_zero_loss(ppo_strategy):
"""A zero-head critic, zero advantages, and zero returns → zero loss."""
strategy, device = ppo_strategy
batch = _make_batch(device=device)
batch["advantages"] = torch.zeros_like(batch["responses"], dtype=torch.float)
batch["returns"] = torch.zeros_like(batch["responses"], dtype=torch.float)
loss = strategy.compute_loss(batch)
assert loss.item() == pytest.approx(0.0, abs=1e-6)
def test_ppo_all_masked_response_tokens_zero_loss(ppo_strategy):
strategy, device = ppo_strategy
batch = _make_batch(device=device)
batch["masks"] = torch.zeros_like(batch["masks"])
loss = strategy.compute_loss(batch)
assert loss.item() == pytest.approx(0.0, abs=1e-6)
def test_ppo_uses_supplied_advantages_without_recomputation(ppo_strategy):
"""Explicit advantages/returns must short-circuit GAE computation."""
strategy, device = ppo_strategy
batch = _make_batch(device=device)
def _fail(*args, **kwargs):
raise AssertionError("advantages were supplied; GAE must not run")
strategy._compute_advantages = _fail
batch["advantages"] = torch.ones_like(batch["responses"], dtype=torch.float)
batch["returns"] = torch.zeros_like(batch["responses"], dtype=torch.float)
loss = strategy.compute_loss(batch)
assert torch.isfinite(loss).item()
def test_ppo_optimizer_step_updates_policy_and_critic(ppo_strategy):
"""optimizer_step steps the policy optimizer and then the critic's."""
strategy, device = ppo_strategy
batch = _make_batch(device=device)
loss = strategy.compute_loss(batch)
loss.backward()
policy_optimizer = torch.optim.SGD(strategy.model.parameters(), lr=0.1)
policy_before = next(strategy.model.parameters()).detach().clone()
critic_before = next(strategy.critic.parameters()).detach().clone()
strategy.optimizer_step(policy_optimizer)
assert not torch.equal(next(strategy.model.parameters()), policy_before)
assert not torch.equal(next(strategy.critic.parameters()), critic_before)
# Critic gradients are cleared after its step.
assert all(p.grad is None for p in strategy.critic.parameters())
def test_ppo_optimizer_step_clips_critic_gradients(ppo_strategy):
strategy, device = ppo_strategy
batch = _make_batch(device=device)
loss = strategy.compute_loss(batch)
loss.backward()
strategy.critic_optimizer = torch.optim.SGD(strategy.critic.parameters(), lr=1.0)
strategy.max_grad_norm = 1e-8
before = next(strategy.critic.parameters()).detach().clone()
strategy.optimizer_step(torch.optim.SGD(strategy.model.parameters(), lr=0.0))
# Clipped-to-zero critic gradients under SGD (no momentum) leave the
# parameters unchanged.
assert torch.equal(next(strategy.critic.parameters()), before)
# ============== prepare_from_rollout / GAE integration ==============
def _make_rollout_result(B=2, G=2, P=6, R=5, device="cpu"):
return RolloutResult(
prompts=torch.randint(3, 200, (B, P), device=device),
prompt_mask=torch.ones(B, P, dtype=torch.bool, device=device),
responses=torch.randint(3, 200, (B, G, R), device=device),
response_mask=torch.ones(B, G, R, dtype=torch.bool, device=device),
rewards=torch.randn(B, G, device=device),
logprobs_old=torch.zeros(B, G, R, device=device),
)
def test_prepare_from_rollout_computes_and_pins_gae(ppo_strategy, monkeypatch):
"""prepare attaches GAE tensors to the result once, then reuses them."""
strategy, device = ppo_strategy
result = _make_rollout_result(device=device)
batch = strategy.prepare_from_rollout(result)
assert result.advantages is not None and result.returns is not None
assert batch["advantages"] is result.advantages
assert batch["returns"] is result.returns
assert batch["advantages"].shape == result.responses.shape
calls = []
original = strategy._compute_advantages
monkeypatch.setattr(
strategy,
"_compute_advantages",
lambda *a, **k: calls.append(1) or original(*a, **k),
)
strategy.prepare_from_rollout(result)
assert not calls, "pinned advantages must not be recomputed on replay"
def test_prepare_from_rollout_respects_response_padding(ppo_strategy):
"""Padded response positions get zero advantages and returns."""
strategy, device = ppo_strategy
result = _make_rollout_result(device=device)
result.response_mask[0, 0, 3:] = False
batch = strategy.prepare_from_rollout(result)
assert torch.count_nonzero(batch["advantages"][0, 0, 3:]) == 0
assert torch.count_nonzero(batch["returns"][0, 0, 3:]) == 0
assert torch.count_nonzero(batch["advantages"][0, 0, :3]) > 0
def test_compute_advantages_matches_hand_computed_gae(ppo_strategy, monkeypatch):
"""_compute_advantages applies terminal rewards and GAE faithfully."""
strategy, device = ppo_strategy
result = _make_rollout_result(B=1, G=1, P=4, R=3, device=device)
result.rewards = torch.tensor([[2.0]], device=device)
result.logprobs_old = torch.zeros(1, 1, 3, device=device)
# ref_model == policy at init → zero KL reward shaping only if the
# policy and ref agree on logprobs; keep ref out of the picture here.
strategy.ref_model = None
fixed_values = torch.tensor([[[0.1, 0.2, 0.3]]], device=device)
monkeypatch.setattr(
strategy_module,
"rollout_token_values",
lambda *args, **kwargs: fixed_values.clone(),
)
advantages, returns = strategy._compute_advantages(
result.prompts,
result.prompt_mask,
result.responses,
result.response_mask,
result.rewards,
result.logprobs_old,
)
rewards = torch.tensor([[[0.0, 0.0, 2.0]]], device=device)
expected_adv, expected_ret = compute_gae(
rewards, fixed_values, result.response_mask, 1.0, 0.95
)
assert torch.allclose(advantages, expected_adv, atol=1e-6)
assert torch.allclose(returns, expected_ret, atol=1e-6)
def test_compute_advantages_folds_kl_penalty_into_rewards(ppo_strategy, monkeypatch):
"""With a ref model, each valid token's reward loses kl_coef·k3."""
strategy, device = ppo_strategy
result = _make_rollout_result(B=1, G=1, P=4, R=2, device=device)
result.rewards = torch.tensor([[1.0]], device=device)
# behaviour policy disagrees with ref by +1 logprob on every token
result.logprobs_old = torch.ones(1, 1, 2, device=device)
fixed_values = torch.zeros(1, 1, 2, device=device)
fixed_ref_logprobs = torch.zeros(1, 1, 2, device=device)
monkeypatch.setattr(
strategy_module,
"rollout_token_values",
lambda *args, **kwargs: fixed_values.clone(),
)
monkeypatch.setattr(
strategy_module,
"rollout_token_logprobs",
lambda *args, **kwargs: {"logprobs": fixed_ref_logprobs.clone()},
)
advantages, _ = strategy._compute_advantages(
result.prompts,
result.prompt_mask,
result.responses,
result.response_mask,
result.rewards,
result.logprobs_old,
)
# per-token reward = -kl_coef·(1 - 0) = -0.01; terminal adds 1.0
expected_rewards = torch.tensor([[-0.01, 0.99]], device=device)
expected_adv, _ = compute_gae(
expected_rewards.unsqueeze(0), fixed_values, result.response_mask, 1.0, 0.95
)
assert torch.allclose(advantages, expected_adv, atol=1e-6)
def test_online_call_returns_finite_loss(ppo_strategy):
strategy, device = ppo_strategy
class _RecordingRunner:
policy_version = 0
def __call__(self, batch):
return _make_rollout_result(device=device), True
def step(self):
pass
def apply_weight_update(self, policy_version, update):
return update()
strategy.set_rollout_runner(_RecordingRunner())
out = strategy({"instruction": ["x"]})
assert torch.isfinite(out["loss"]).item()
assert "policy_loss" in out["metrics"]
assert "value_loss" in out["metrics"]
assert "explained_variance" in out["metrics"]