ViperEkura
350e4a1849
- 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
56 lines
2.2 KiB
Python
56 lines
2.2 KiB
Python
"""Value (critic) model for actor-critic RL training."""
|
|
|
|
from typing import Dict, Optional
|
|
|
|
import torch.nn as nn
|
|
from torch import Tensor
|
|
|
|
from astrai.config.model_config import AutoRegressiveLMConfig
|
|
from astrai.model.automodel import ModelFactory
|
|
from astrai.model.components.linear import Linear
|
|
from astrai.model.transformer import AutoRegressiveLM, process_attention_mask
|
|
|
|
|
|
@ModelFactory.register("value_model")
|
|
class ValueModel(AutoRegressiveLM):
|
|
"""Critic scoring each state with a scalar instead of vocab logits.
|
|
|
|
Inherits the ``AutoRegressiveLM`` components so a policy checkpoint can
|
|
warm-start the critic backbone (``load_state_dict(..., strict=False)``);
|
|
only ``value_head`` keeps its fresh initialization. The inherited
|
|
``lm_head`` parameters stay dormant — the forward below never projects
|
|
through them — so checkpoints round-trip with stable keys. The trunk
|
|
pass mirrors ``AutoRegressiveLM.forward`` for training-style input;
|
|
``tests/trainer/test_ppo_strategy.py`` pins the two to identical
|
|
hidden states.
|
|
"""
|
|
|
|
def __init__(self, config: AutoRegressiveLMConfig):
|
|
super().__init__(config)
|
|
self.value_head = Linear(config.hidden_size, 1, bias=True)
|
|
# Zero head so training starts from V(s) == 0 and the first GAE
|
|
# advantages are driven purely by rewards.
|
|
nn.init.zeros_(self.value_head.weight)
|
|
nn.init.zeros_(self.value_head.bias)
|
|
|
|
def forward(
|
|
self,
|
|
input_ids: Tensor,
|
|
input_mask: Optional[Tensor] = None,
|
|
position_ids: Optional[Tensor] = None,
|
|
) -> Dict[str, Tensor]:
|
|
if input_ids.ndim != 2:
|
|
raise ValueError("critic input_ids must be [batch, seq_len]")
|
|
x = self.embed_tokens(input_ids)
|
|
rotary_emb = self.rotary_embedding(x, position_ids)
|
|
attn_mask = process_attention_mask(input_mask)
|
|
use_sdpa_causal_mask = attn_mask is None
|
|
|
|
for layer in self.layers:
|
|
x = layer(x, rotary_emb, attn_mask, None, use_sdpa_causal_mask, None)[
|
|
"hidden_states"
|
|
]
|
|
hidden_states = self.norm(x)
|
|
values = self.value_head(hidden_states).squeeze(-1)
|
|
return {"values": values}
|