Spaces:
Running
Running
Add frozen-critic PPO and show objective gradients directly on tokens
Browse filesUse the pinned ORZ 0.5B value critic with per-token GAE and editable values. Replace toy probability-change chips with exact objective gradients, keep collected probabilities fixed, and explain the optimizer/forward-pass distinction. Preserve monochrome RL modes and fix mobile navigation overflow.
- .dockerignore +2 -0
- Dockerfile +2 -2
- README.md +24 -6
- critic.py +105 -0
- grpo.py +65 -38
- ppo.py +126 -0
- preload_models.py +3 -1
- server.py +58 -4
- static/grpo.css +11 -3
- static/grpo.html +11 -11
- static/grpo.js +62 -42
- static/index.html +1 -1
- static/ppo.html +126 -0
- static/styles.css +2 -0
.dockerignore
CHANGED
|
@@ -4,6 +4,8 @@
|
|
| 4 |
!runtime.py
|
| 5 |
!server.py
|
| 6 |
!grpo.py
|
|
|
|
|
|
|
| 7 |
!preload_models.py
|
| 8 |
!static/
|
| 9 |
!static/**
|
|
|
|
| 4 |
!runtime.py
|
| 5 |
!server.py
|
| 6 |
!grpo.py
|
| 7 |
+
!ppo.py
|
| 8 |
+
!critic.py
|
| 9 |
!preload_models.py
|
| 10 |
!static/
|
| 11 |
!static/**
|
Dockerfile
CHANGED
|
@@ -12,14 +12,14 @@ RUN python -c 'import subprocess, sys, tomllib; deps = tomllib.load(open("/tmp/p
|
|
| 12 |
|
| 13 |
RUN useradd --create-home --uid 1000 user
|
| 14 |
WORKDIR /app
|
| 15 |
-
COPY --chown=user:user runtime.py preload_models.py ./
|
| 16 |
USER user
|
| 17 |
|
| 18 |
# Cache only the pinned inference files, excluding training checkpoints.
|
| 19 |
RUN python preload_models.py
|
| 20 |
ENV HF_HUB_OFFLINE=1
|
| 21 |
|
| 22 |
-
COPY --chown=user:user server.py grpo.py ./
|
| 23 |
COPY --chown=user:user static ./static
|
| 24 |
|
| 25 |
EXPOSE 7860
|
|
|
|
| 12 |
|
| 13 |
RUN useradd --create-home --uid 1000 user
|
| 14 |
WORKDIR /app
|
| 15 |
+
COPY --chown=user:user runtime.py critic.py preload_models.py ./
|
| 16 |
USER user
|
| 17 |
|
| 18 |
# Cache only the pinned inference files, excluding training checkpoints.
|
| 19 |
RUN python preload_models.py
|
| 20 |
ENV HF_HUB_OFFLINE=1
|
| 21 |
|
| 22 |
+
COPY --chown=user:user server.py grpo.py ppo.py ./
|
| 23 |
COPY --chown=user:user static ./static
|
| 24 |
|
| 25 |
EXPOSE 7860
|
README.md
CHANGED
|
@@ -19,7 +19,7 @@ tags:
|
|
| 19 |
|
| 20 |
# Layer lab — Qwen & Gemma
|
| 21 |
|
| 22 |
-
A website for inspecting **Qwen3-0.6B** and **Gemma 3 270M TinyGSM reasoning**: next-token predictions at every position after every decoder layer, attention from every query head, and
|
| 23 |
|
| 24 |
| Checkpoint | Layers | Query / KV heads per layer | Reasoning |
|
| 25 |
| --- | ---: | ---: | --- |
|
|
@@ -33,7 +33,7 @@ cd qwen-lab
|
|
| 33 |
uv run --python 3.11 python server.py --port 8043
|
| 34 |
```
|
| 35 |
|
| 36 |
-
Open [http://localhost:8043](http://localhost:8043). If the dependencies are already installed, `python3.11 server.py` is sufficient. The first launch downloads both checkpoints and
|
| 37 |
|
| 38 |
The local server binds to `127.0.0.1` by default. Computation runs on the server hosting the website: your computer for a local launch, or the Space for a hosted launch. Prompts are not sent to an inference API. On first local launch, the library downloads checkpoint files from Hugging Face. Compatible Macs use MPS automatically; other systems use CPU. To force CPU, run `QWEN_DEVICE=cpu uv run --python 3.11 python server.py`. Computation uses float32 from the original BF16 checkpoint, without quantization.
|
| 39 |
|
|
@@ -41,7 +41,7 @@ The local server binds to `127.0.0.1` by default. Computation runs on the server
|
|
| 41 |
|
| 42 |
The repository is ready for a **Docker Space** on **CPU Basic** (2 vCPU, 16 GB RAM). CPU Basic has no hourly hardware charge; creating Docker Spaces currently requires PRO for personal accounts or Team/Enterprise for organizations. CPU inference is slower than MPS, especially for long contexts.
|
| 43 |
|
| 44 |
-
Create a Docker Space and upload `README.md`, `Dockerfile`, `.dockerignore`, `pyproject.toml`, `preload_models.py`, `runtime.py`, `grpo.py`, `server.py`, and `static/` to the repository root. The Docker image installs CPU-only PyTorch and downloads both pinned checkpoints during the build, excluding training checkpoints and optimizer state. Runtime loads the cached files offline. No model token or external inference service is required. The server listens on `0.0.0.0:7860` as required by Spaces.
|
| 45 |
|
| 46 |
The app uses a single Python process. Results live in memory, with the two most recent analyses retained per model across all visitors. A Space restart or sleep clears them; an evicted result needs a new analysis. CPU Basic can sleep after inactivity and wake when visited. No persistent storage is required.
|
| 47 |
|
|
@@ -68,12 +68,28 @@ The notation follows the maximized objective **𝒥_GRPO**: `r_i` is a rollout r
|
|
| 68 |
|
| 69 |
At the collection policy, `ρ_i,t = 1`, `K_i,t = 0`, and neither clipping nor KL changes the first gradient. **Candidate drift** ranges from **−3 to +3**, with zero in the middle, and sets `logit(π_θ) = logit(π_old) + d Â_i`. Positive drift moves with the advantage; negative drift moves against it. Clipping still allows gradients that recover from unfavorable policy moves even outside the clip interval. Here `logit(π)` means log-odds of the sampled token against all other vocabulary entries. This drift is hypothetical, not another model forward pass.
|
| 70 |
|
| 71 |
-
|
| 72 |
|
| 73 |
-
**
|
|
|
|
|
|
|
| 74 |
|
| 75 |
The API uses `POST /api/grpo/runs` to start sampling, `GET /api/grpo/runs/{id}` for progress and samples, `POST /api/grpo/runs/{id}/cancel` to stop, and `POST /api/grpo/runs/{id}/preview` for fast reward/objective recalculation. Sampling is asynchronous to keep long CPU requests out of the web proxy's request timeout.
|
| 76 |
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 77 |
## Interface
|
| 78 |
|
| 79 |
- **Model**: switch between Qwen and Gemma. Shapes, token IDs, templates, GQA mapping, method notes and checkpoint links follow the selected model. A switch clears previous results while preserving the input. Both models stay loaded; switching does not reload weights.
|
|
@@ -127,7 +143,7 @@ The interface limit is **512 tokens including the continuation**, independently
|
|
| 127 |
## Validation
|
| 128 |
|
| 129 |
```sh
|
| 130 |
-
uv run --python 3.11 --group dev python -m pytest tests/test_model.py tests/test_grpo.py -q
|
| 131 |
|
| 132 |
# Keep the Python server running in another terminal:
|
| 133 |
npm install
|
|
@@ -140,4 +156,6 @@ Numerical integration checks use both actual pinned checkpoints: the official re
|
|
| 140 |
|
| 141 |
GRPO checks compare analytic logit gradients with finite differences of the complete objective, verify softmax renormalization and clipping/KL behavior, compare real sampled-token probabilities against an independent teacher-forced forward pass in both checkpoints, check seeded sampling and cancellation, and verify unchanged parameter versions, probe values, absent parameter gradients, and evaluation mode. Browser checks cover both models, editable rewards, zero-advantage groups, ablations, prefixes and thinking, export, cancellation, and responsive layouts.
|
| 142 |
|
|
|
|
|
|
|
| 143 |
Browser checks cover the append column, predicted and manual additions in both models, tokenizer previews, limit errors, unchanged prefix IDs, token sources and the active attention color mode after appending. They also cover the reversed layer stack, the sticky bottom input row, native resizing, rendered colors for every head, exact weights and masks at first, intermediate and last query positions, automatic column visibility, peak navigation, linear and logarithmic colors, model switching, templates, control tokens, reasoning, exports, and the absence of page overflow at widths from 320 to 1440 px. Screenshots are saved in `test-results/`.
|
|
|
|
| 19 |
|
| 20 |
# Layer lab — Qwen & Gemma
|
| 21 |
|
| 22 |
+
A website for inspecting **Qwen3-0.6B** and **Gemma 3 270M TinyGSM reasoning**: next-token predictions at every position after every decoder layer, attention from every query head, and frozen-model GRPO/PPO sandboxes. PPO uses a separate pretrained ORZ 0.5B value critic. Run it locally or on Hugging Face Spaces. Qwen remains the default actor.
|
| 23 |
|
| 24 |
| Checkpoint | Layers | Query / KV heads per layer | Reasoning |
|
| 25 |
| --- | ---: | ---: | --- |
|
|
|
|
| 33 |
uv run --python 3.11 python server.py --port 8043
|
| 34 |
```
|
| 35 |
|
| 36 |
+
Open [http://localhost:8043](http://localhost:8043). If the dependencies are already installed, `python3.11 server.py` is sufficient. The first launch downloads both actor checkpoints and the value critic from Hugging Face; subsequent launches use the standard Hugging Face cache. No API key is required. Qwen, Gemma, and the critic load in that order in the background. Each actor becomes available after loading; PPO also waits for the critic. The layer explorer's initial example runs automatically when the model is ready.
|
| 37 |
|
| 38 |
The local server binds to `127.0.0.1` by default. Computation runs on the server hosting the website: your computer for a local launch, or the Space for a hosted launch. Prompts are not sent to an inference API. On first local launch, the library downloads checkpoint files from Hugging Face. Compatible Macs use MPS automatically; other systems use CPU. To force CPU, run `QWEN_DEVICE=cpu uv run --python 3.11 python server.py`. Computation uses float32 from the original BF16 checkpoint, without quantization.
|
| 39 |
|
|
|
|
| 41 |
|
| 42 |
The repository is ready for a **Docker Space** on **CPU Basic** (2 vCPU, 16 GB RAM). CPU Basic has no hourly hardware charge; creating Docker Spaces currently requires PRO for personal accounts or Team/Enterprise for organizations. CPU inference is slower than MPS, especially for long contexts.
|
| 43 |
|
| 44 |
+
Create a Docker Space and upload `README.md`, `Dockerfile`, `.dockerignore`, `pyproject.toml`, `preload_models.py`, `runtime.py`, `critic.py`, `grpo.py`, `ppo.py`, `server.py`, and `static/` to the repository root. The Docker image installs CPU-only PyTorch and downloads both pinned actor checkpoints and the pinned critic during the build, excluding training checkpoints and optimizer state. Runtime loads the cached files offline. No model token or external inference service is required. The server listens on `0.0.0.0:7860` as required by Spaces.
|
| 45 |
|
| 46 |
The app uses a single Python process. Results live in memory, with the two most recent analyses retained per model across all visitors. A Space restart or sleep clears them; an evicted result needs a new analysis. CPU Basic can sleep after inactivity and wake when visited. No persistent storage is required.
|
| 47 |
|
|
|
|
| 68 |
|
| 69 |
At the collection policy, `ρ_i,t = 1`, `K_i,t = 0`, and neither clipping nor KL changes the first gradient. **Candidate drift** ranges from **−3 to +3**, with zero in the middle, and sets `logit(π_θ) = logit(π_old) + d Â_i`. Positive drift moves with the advantage; negative drift moves against it. Clipping still allows gradients that recover from unfavorable policy moves even outside the clip interval. Here `logit(π)` means log-odds of the sampled token against all other vocabulary entries. This drift is hypothetical, not another model forward pass.
|
| 70 |
|
| 71 |
+
Every chip automatically shows its collected probability and **`g_i,t = ∂𝒥_GRPO/∂log π_θ(y_i,t | h_i,t)`**, including group/token averaging. Gradients recalculate immediately after editing rewards or objective settings. Positive g locally favors increasing this log-probability; negative g favors decreasing it. Selecting a token shows its unaveraged objective/loss, policy and KL gradient contributions, candidate probability, ratio, and clipping branch. Both sandboxes use flat monochrome styling.
|
| 72 |
|
| 73 |
+
**A gradient is not a probability change.** Actual training computes `∇θ𝒥 = Σ g_i,t ∇θ log π_θ`, updates shared weights with an optimizer, and runs another forward pass to obtain new probabilities. Shared parameters couple tokens; the optimizer and its state affect the result. This sandbox computes no parameter gradients, optimizer updates, or updated-model probabilities. Clipping removes a policy-gradient incentive; it does not clamp probabilities. JSON exports contain samples, rewards, settings, objective values, and local derivatives. For API compatibility, `loss = -objective`, loss derivatives, and legacy independent-logit toy `after`/`delta` fields remain available. The website requests a zero toy step and never displays those fields as measured or predicted training changes.
|
| 74 |
+
|
| 75 |
+
Both checkpoints stay unchanged. Up to two completed groups per actor and algorithm are retained alongside an active group; older groups and restarts require resampling. Generation shares the actor's existing computation lock with the layer explorer.
|
| 76 |
|
| 77 |
The API uses `POST /api/grpo/runs` to start sampling, `GET /api/grpo/runs/{id}` for progress and samples, `POST /api/grpo/runs/{id}/cancel` to stop, and `POST /api/grpo/runs/{id}/preview` for fast reward/objective recalculation. Sampling is asynchronous to keep long CPU requests out of the web proxy's request timeout.
|
| 78 |
|
| 79 |
+
## PPO sandbox
|
| 80 |
+
|
| 81 |
+
Open `/ppo`. The same real Qwen/Gemma sampling, templates, reasoning controls, rewards, clipping, KL, signed illustrative drift, and exports are available. Five rollouts remain the default. Each token additionally shows **V(history before the token)** and its **GAE advantage**, alongside the collected probability and actor-objective gradient g. The inspector exposes the immediate reward, next V, TD residual, raw/normalized advantage, detached λ-return, and value loss. Values can be overridden per state, reset to checkpoint predictions, or ablated to zero.
|
| 82 |
+
|
| 83 |
+
The critic is [Open-Reasoner-Zero/Open-Reasoner-Zero-Critic-0.5B](https://huggingface.co/Open-Reasoner-Zero/Open-Reasoner-Zero-Critic-0.5B), pinned to `8c1c681b55b641811b88a41d9a81017f90ac6369` (MIT). It contains a trained Qwen2.5 backbone and scalar `value_head.weight`, loaded strictly using local Transformers classes without executing remote code. Its optional mean/std normalization follows the [original critic implementation](https://github.com/Open-Reasoner-Zero/Open-Reasoner-Zero/blob/main/orz/ppo/models.py). No random head or reward-model substitution is used. A scalar value model has no sampling temperature.
|
| 84 |
+
|
| 85 |
+
**The critic was trained on mathematical reasoning with another policy and reward distribution.** Its estimates are illustrative, not calibrated values for Qwen3/Gemma or arbitrary custom targets. Changing an answer or reward does not retrain V. The interface explains this mismatch and permits manual value experiments. Both actor and critic stay frozen.
|
| 86 |
+
|
| 87 |
+
Critic histories are the exact decoded actor prefixes before each action, including control tokens, independently encoded with the critic's own tokenizer. Future text is never used to find a token boundary. Incremental inference crops the critic's KV cache to the longest unchanged token prefix when BPE retokenization changes a suffix. This also handles Gemma's different tokenizer and partial Unicode pieces. A critic history exceeding 4096 critic tokens is rejected without truncation.
|
| 88 |
+
|
| 89 |
+
EOS **and the token budget** terminate the finite-horizon episode. The final action receives the rollout score; earlier rewards and terminal V are zero. A token-budget completion is scored as displayed, with no infinite-horizon bootstrap. For each rollout, `δ_t = r_t + γ V(h_{t+1}) − V(h_t)` and `Â_t = δ_t + γλ Â_{t+1}` are evaluated backward. Defaults are γ = 1 and λ = 0.95. Optional actor-advantage normalization uses all generated tokens, sample SD + 10⁻⁸, and is off by default. The unnormalized detached λ-return is `R̂_t = Â_t + V(h_t)`; critic loss is `½(V(h_t) − stopgrad(R̂_t))²`. The critic loss and actor objective both average each rollout's token mean, then average the group. KL is an actor-objective penalty, not an additional reward in GAE. Actor and critic parameter gradients are not computed.
|
| 90 |
+
|
| 91 |
+
PPO routes mirror GRPO under `/api/ppo/runs`. A run progresses through `sampling`, `valuing`, and `ready`; cancellation works in both computation phases. The preview accepts `gamma`, `gae_lambda`, `critic`, and optional per-token `values` overrides in addition to the shared objective settings. Overrides and reward edits reuse fixed samples and cached critic predictions. `/api/status` also reports critic loading status and availability.
|
| 92 |
+
|
| 93 |
## Interface
|
| 94 |
|
| 95 |
- **Model**: switch between Qwen and Gemma. Shapes, token IDs, templates, GQA mapping, method notes and checkpoint links follow the selected model. A switch clears previous results while preserving the input. Both models stay loaded; switching does not reload weights.
|
|
|
|
| 143 |
## Validation
|
| 144 |
|
| 145 |
```sh
|
| 146 |
+
uv run --python 3.11 --group dev python -m pytest tests/test_model.py tests/test_grpo.py tests/test_ppo.py -q
|
| 147 |
|
| 148 |
# Keep the Python server running in another terminal:
|
| 149 |
npm install
|
|
|
|
| 156 |
|
| 157 |
GRPO checks compare analytic logit gradients with finite differences of the complete objective, verify softmax renormalization and clipping/KL behavior, compare real sampled-token probabilities against an independent teacher-forced forward pass in both checkpoints, check seeded sampling and cancellation, and verify unchanged parameter versions, probe values, absent parameter gradients, and evaluation mode. Browser checks cover both models, editable rewards, zero-advantage groups, ablations, prefixes and thinking, export, cancellation, and responsive layouts.
|
| 158 |
|
| 159 |
+
PPO checks cover GAE/TD recursion, discounted returns, finite-horizon terminal values, detached λ-return targets, value overrides and ablations, global advantage normalization, actor and critic derivatives against finite differences, and real actor/critic jobs with both actor tokenizers. Cached critic values are compared to separate full-history forward passes, including BPE boundary changes and incomplete Unicode pieces, and both networks remain unchanged. `npm run test:ppo` checks the real browser flow, immediate token gradients/V/GAE, value edits, ablations, JSON export, monochrome colors, and responsive 320–1440px layouts.
|
| 160 |
+
|
| 161 |
Browser checks cover the append column, predicted and manual additions in both models, tokenizer previews, limit errors, unchanged prefix IDs, token sources and the active attention color mode after appending. They also cover the reversed layer stack, the sticky bottom input row, native resizing, rendered colors for every head, exact weights and masks at first, intermediate and last query positions, automatic column visibility, peak navigation, linear and logarithmic colors, model switching, templates, control tokens, reasoning, exports, and the absence of page overflow at widths from 320 to 1440 px. Screenshots are saved in `test-results/`.
|
critic.py
ADDED
|
@@ -0,0 +1,105 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
"""Pinned, frozen ORZ state-value critic; no remote Python code is executed."""
|
| 2 |
+
from __future__ import annotations
|
| 3 |
+
|
| 4 |
+
import math
|
| 5 |
+
import os
|
| 6 |
+
import threading
|
| 7 |
+
|
| 8 |
+
import torch
|
| 9 |
+
from torch import nn
|
| 10 |
+
from transformers import AutoConfig, AutoTokenizer, Qwen2Model, Qwen2PreTrainedModel
|
| 11 |
+
|
| 12 |
+
CRITIC_ID = "Open-Reasoner-Zero/Open-Reasoner-Zero-Critic-0.5B"
|
| 13 |
+
CRITIC_REVISION = "8c1c681b55b641811b88a41d9a81017f90ac6369"
|
| 14 |
+
CRITIC_MAX_TOKENS = 4096
|
| 15 |
+
|
| 16 |
+
|
| 17 |
+
class ValueModel(Qwen2PreTrainedModel):
|
| 18 |
+
"""The checkpoint's Qwen2 backbone and trained scalar value_head."""
|
| 19 |
+
|
| 20 |
+
def __init__(self, config):
|
| 21 |
+
super().__init__(config)
|
| 22 |
+
self.model = Qwen2Model(config)
|
| 23 |
+
self.value_head = nn.Linear(config.hidden_size, 1, bias=False)
|
| 24 |
+
self.post_init()
|
| 25 |
+
|
| 26 |
+
|
| 27 |
+
class CriticRuntime:
|
| 28 |
+
def __init__(self):
|
| 29 |
+
self.phase = "loading"
|
| 30 |
+
self.message = "Loading the ORZ 0.5B value critic…"
|
| 31 |
+
self.model = self.tokenizer = None
|
| 32 |
+
self.lock = threading.Lock()
|
| 33 |
+
self.device = "cpu"
|
| 34 |
+
|
| 35 |
+
def status(self):
|
| 36 |
+
return dict(phase=self.phase, message=self.message, model=CRITIC_ID,
|
| 37 |
+
revision=CRITIC_REVISION, device=self.device, busy=self.lock.locked())
|
| 38 |
+
|
| 39 |
+
def load(self):
|
| 40 |
+
try:
|
| 41 |
+
config = AutoConfig.from_pretrained(CRITIC_ID, revision=CRITIC_REVISION)
|
| 42 |
+
self.tokenizer = AutoTokenizer.from_pretrained(CRITIC_ID, revision=CRITIC_REVISION)
|
| 43 |
+
self.model, info = ValueModel.from_pretrained(
|
| 44 |
+
CRITIC_ID, revision=CRITIC_REVISION, config=config, dtype=torch.float32,
|
| 45 |
+
attn_implementation="eager", output_loading_info=True)
|
| 46 |
+
if any(info.get(key) for key in ("missing_keys", "unexpected_keys", "mismatched_keys", "error_msgs")):
|
| 47 |
+
raise ValueError(f"The value checkpoint did not load exactly: {info}")
|
| 48 |
+
requested = os.environ.get("QWEN_DEVICE", "auto")
|
| 49 |
+
self.device = ("mps" if torch.backends.mps.is_available() else "cpu") if requested == "auto" else requested
|
| 50 |
+
self.model.to(self.device).eval().requires_grad_(False)
|
| 51 |
+
self.mean = float(getattr(config, "mean", 0.0))
|
| 52 |
+
self.std = float(getattr(config, "std", 1.0))
|
| 53 |
+
if not math.isfinite(self.std) or self.std <= 0:
|
| 54 |
+
raise ValueError("The critic normalization standard deviation must be positive.")
|
| 55 |
+
self.normalize = bool(getattr(config, "normalize_reward", False))
|
| 56 |
+
self.phase = "ready"
|
| 57 |
+
self.message = f"ORZ 0.5B · trained value-head · {self.device.upper()} · frozen"
|
| 58 |
+
except Exception as exc:
|
| 59 |
+
import traceback
|
| 60 |
+
traceback.print_exc()
|
| 61 |
+
self.phase = "error"
|
| 62 |
+
self.message = f"Could not load the value critic: {exc}"
|
| 63 |
+
|
| 64 |
+
@torch.inference_mode()
|
| 65 |
+
def values(self, policy, prefix_ids, tokens, progress, cancelled):
|
| 66 |
+
"""Evaluate V(history BEFORE each action), using the critic's tokenizer.
|
| 67 |
+
|
| 68 |
+
Retokenize each exact policy history independently. Its final BPE pieces
|
| 69 |
+
can change as a character is appended, so crop KV cache to the longest
|
| 70 |
+
unchanged prefix. Never align a state using future text or actor IDs.
|
| 71 |
+
"""
|
| 72 |
+
result, previous, previous_values, cache = [], [], [], None
|
| 73 |
+
sampled = [token["id"] for token in tokens]
|
| 74 |
+
for position in range(len(tokens)):
|
| 75 |
+
if cancelled():
|
| 76 |
+
return None
|
| 77 |
+
progress(position)
|
| 78 |
+
history = policy.tokenizer.decode(prefix_ids + sampled[:position],
|
| 79 |
+
skip_special_tokens=False, clean_up_tokenization_spaces=False)
|
| 80 |
+
ids = self.tokenizer.encode(history, add_special_tokens=False)
|
| 81 |
+
if not ids or len(ids) > CRITIC_MAX_TOKENS:
|
| 82 |
+
raise ValueError(f"A critic history must contain 1–{CRITIC_MAX_TOKENS} critic tokens. No text was truncated.")
|
| 83 |
+
common = 0
|
| 84 |
+
for old, new in zip(previous, ids):
|
| 85 |
+
if old != new:
|
| 86 |
+
break
|
| 87 |
+
common += 1
|
| 88 |
+
if cache is not None:
|
| 89 |
+
cache.crop(common)
|
| 90 |
+
state_values = previous_values[:common]
|
| 91 |
+
if common < len(ids):
|
| 92 |
+
output = self.model.model(torch.tensor([ids[common:]], device=self.device),
|
| 93 |
+
past_key_values=cache, use_cache=True, output_attentions=False, output_hidden_states=False)
|
| 94 |
+
cache = output.past_key_values
|
| 95 |
+
scores = self.model.value_head(output.last_hidden_state).squeeze(-1).float()
|
| 96 |
+
if self.normalize:
|
| 97 |
+
scores = (scores - self.mean) / self.std
|
| 98 |
+
state_values += scores[0].cpu().tolist()
|
| 99 |
+
del output, scores
|
| 100 |
+
value = state_values[-1]
|
| 101 |
+
if not math.isfinite(value):
|
| 102 |
+
raise ValueError("The critic returned a nonfinite value.")
|
| 103 |
+
result.append(value)
|
| 104 |
+
previous, previous_values = ids, state_values
|
| 105 |
+
return result
|
grpo.py
CHANGED
|
@@ -1,4 +1,4 @@
|
|
| 1 |
-
"""Frozen-policy rollouts and
|
| 2 |
from __future__ import annotations
|
| 3 |
|
| 4 |
import copy
|
|
@@ -39,6 +39,46 @@ def sigmoid(x):
|
|
| 39 |
return e / (1 + e)
|
| 40 |
|
| 41 |
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 42 |
def preview_step(rollouts, payload):
|
| 43 |
"""Differentiate the GRPO loss, then step each occurrence's sampled-token logit.
|
| 44 |
|
|
@@ -68,49 +108,25 @@ def preview_step(rollouts, payload):
|
|
| 68 |
raise ValueError("A rollout has no sampled tokens.")
|
| 69 |
weight = 1 / (count * length)
|
| 70 |
for token in row["tokens"]:
|
| 71 |
-
|
| 72 |
-
|
| 73 |
-
|
| 74 |
-
|
| 75 |
-
clipped_ratio = min(max(ratio, 1 - payload.epsilon), 1 + payload.epsilon)
|
| 76 |
-
# At the boundary, use the unclipped branch's one-sided derivative.
|
| 77 |
-
clipped = payload.clipping and ((advantage > 0 and ratio > 1 + payload.epsilon)
|
| 78 |
-
or (advantage < 0 and ratio < 1 - payload.epsilon))
|
| 79 |
-
surrogate = min(ratio * advantage, clipped_ratio * advantage) if payload.clipping else ratio * advantage
|
| 80 |
-
reference_ratio = old / current
|
| 81 |
-
kl = max(0.0, reference_ratio - math.log(reference_ratio) - 1)
|
| 82 |
-
beta = payload.beta if payload.kl else 0.0
|
| 83 |
-
# d loss / d log(pi(a|s)); KL uses the sampled k3 estimator.
|
| 84 |
-
policy_gradient = 0.0 if clipped else -advantage * ratio
|
| 85 |
-
kl_gradient = beta * (1 - reference_ratio)
|
| 86 |
-
logp_gradient = policy_gradient + kl_gradient
|
| 87 |
-
logit_gradient = weight * logp_gradient * (1 - current)
|
| 88 |
-
after = sigmoid(math.log(current) - math.log1p(-current) - payload.step_size * logit_gradient)
|
| 89 |
-
loss = -surrogate + beta * kl
|
| 90 |
-
tokens.append(dict(old=old, reference=old, current=current, after=after,
|
| 91 |
-
delta=after-old, step_delta=after-current,
|
| 92 |
-
ratio=ratio, clipped_ratio=clipped_ratio, clipped=clipped,
|
| 93 |
-
kl=kl, policy_gradient=policy_gradient, kl_gradient=kl_gradient,
|
| 94 |
-
logp_gradient=logp_gradient, logit_gradient=logit_gradient,
|
| 95 |
-
weight=weight, surrogate=surrogate, loss=loss,
|
| 96 |
-
objective=-loss, objective_policy_gradient=-policy_gradient,
|
| 97 |
-
objective_kl_gradient=-kl_gradient,
|
| 98 |
-
objective_logp_gradient=-logp_gradient,
|
| 99 |
-
objective_logit_gradient=-logit_gradient))
|
| 100 |
-
total_loss += weight * loss
|
| 101 |
-
total_kl += weight * kl
|
| 102 |
token_count += 1
|
| 103 |
-
clipped_count += int(clipped)
|
| 104 |
rows.append(dict(reward=reward, manual=override is not None, advantage=advantage,
|
| 105 |
scored_answer=scored_answer(row["text"]), tokens=tokens))
|
| 106 |
# J_GRPO is maximized; retain the equivalent loss L = -J for API compatibility.
|
| 107 |
return dict(mean=mean, std=std, objective=-total_loss, loss=total_loss, mean_kl=total_kl,
|
| 108 |
clipped_count=clipped_count, token_count=token_count,
|
| 109 |
all_equal=std == 0, weights_updated=False,
|
|
|
|
| 110 |
probability_space="independent sampled-token logits", rollouts=rows)
|
| 111 |
|
| 112 |
|
| 113 |
class GRPOJobs:
|
|
|
|
|
|
|
| 114 |
def __init__(self):
|
| 115 |
self.jobs = OrderedDict()
|
| 116 |
self.lock = threading.RLock()
|
|
@@ -140,6 +156,7 @@ class GRPOJobs:
|
|
| 140 |
raise BlockingIOError("This model is already computing. Try again when it finishes.")
|
| 141 |
sid = uuid.uuid4().hex
|
| 142 |
job = dict(id=sid, phase="sampling", message="Starting the rollout group…",
|
|
|
|
| 143 |
context=context, prefix_tokens=[runtime.token(t) for t in ids],
|
| 144 |
config=payload.model_dump(), rollouts=[], elapsed=0.0,
|
| 145 |
weights_updated=False, **runtime.metadata(), _cancel=threading.Event())
|
|
@@ -147,7 +164,7 @@ class GRPOJobs:
|
|
| 147 |
with self.lock:
|
| 148 |
self.jobs[sid] = job
|
| 149 |
completed = [key for key, value in self.jobs.items()
|
| 150 |
-
if value["model"] == runtime.model_id and value["phase"]
|
| 151 |
for key in completed[:-2]:
|
| 152 |
del self.jobs[key]
|
| 153 |
threading.Thread(target=self._sample, args=(runtime, job, ids), daemon=True).start()
|
|
@@ -156,12 +173,18 @@ class GRPOJobs:
|
|
| 156 |
raise
|
| 157 |
return self.snapshot(sid)
|
| 158 |
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 159 |
@torch.inference_mode()
|
| 160 |
def _sample(self, runtime, job, prefix_ids):
|
| 161 |
started = time.perf_counter()
|
| 162 |
config = job["config"]
|
| 163 |
phase = "ready"
|
| 164 |
-
final_message = "Group complete. Assign rewards and inspect the
|
| 165 |
try:
|
| 166 |
generator = torch.Generator(device="cpu").manual_seed(config["seed"])
|
| 167 |
eos = runtime.model.generation_config.eos_token_id
|
|
@@ -179,7 +202,7 @@ class GRPOJobs:
|
|
| 179 |
final_message = "Sampling stopped. Model weights are unchanged."
|
| 180 |
return
|
| 181 |
message = f"Rollout {index+1} / {config['group_size']} · token {position+1} / {config['max_new_tokens']}"
|
| 182 |
-
runtime.active = f"
|
| 183 |
with self.lock:
|
| 184 |
job.update(message=message, elapsed=round(time.perf_counter()-started, 2))
|
| 185 |
output = runtime.model(ids, past_key_values=cache, use_cache=True,
|
|
@@ -202,17 +225,21 @@ class GRPOJobs:
|
|
| 202 |
del cache
|
| 203 |
with self.lock:
|
| 204 |
row.update(finished=True, stop_reason="eos" if sampled[-1] in eos else "length")
|
|
|
|
|
|
|
|
|
|
| 205 |
except Exception as exc:
|
| 206 |
import traceback
|
| 207 |
traceback.print_exc()
|
| 208 |
phase = "error"
|
| 209 |
-
final_message = f"
|
| 210 |
finally:
|
| 211 |
with self.lock:
|
| 212 |
job.update(phase=phase, message=final_message, elapsed=round(time.perf_counter()-started, 2))
|
| 213 |
runtime.active = None
|
| 214 |
runtime.lock.release()
|
|
|
|
| 215 |
completed = [key for key, value in self.jobs.items()
|
| 216 |
-
if value["model"] == runtime.model_id and value["phase"]
|
| 217 |
for key in completed[:-2]:
|
| 218 |
del self.jobs[key]
|
|
|
|
| 1 |
+
"""Frozen-policy rollouts and exact local GRPO objective derivatives."""
|
| 2 |
from __future__ import annotations
|
| 3 |
|
| 4 |
import copy
|
|
|
|
| 39 |
return e / (1 + e)
|
| 40 |
|
| 41 |
|
| 42 |
+
def token_objective(token, advantage, weight, payload):
|
| 43 |
+
"""The clipped actor objective and its exact local derivatives.
|
| 44 |
+
|
| 45 |
+
Advantages, reference probabilities, and collection probabilities are
|
| 46 |
+
detached. The derivative in log-probability coordinates is a coefficient
|
| 47 |
+
for backpropagation, not a prediction of a parameter update's effect.
|
| 48 |
+
Legacy independent-logit step fields remain an explicitly toy calculation.
|
| 49 |
+
"""
|
| 50 |
+
old = min(max(token["probability"], 1e-12), 1 - 1e-12)
|
| 51 |
+
old_logit = math.log(old) - math.log1p(-old)
|
| 52 |
+
current = old if payload.drift == 0 else min(max(sigmoid(old_logit + payload.drift * advantage), 1e-12), 1 - 1e-12)
|
| 53 |
+
ratio = current / old
|
| 54 |
+
clipped_ratio = min(max(ratio, 1 - payload.epsilon), 1 + payload.epsilon)
|
| 55 |
+
clipped = payload.clipping and ((advantage > 0 and ratio > 1 + payload.epsilon)
|
| 56 |
+
or (advantage < 0 and ratio < 1 - payload.epsilon))
|
| 57 |
+
surrogate = min(ratio * advantage, clipped_ratio * advantage) if payload.clipping else ratio * advantage
|
| 58 |
+
reference_ratio = old / current
|
| 59 |
+
kl = max(0.0, reference_ratio - math.log(reference_ratio) - 1)
|
| 60 |
+
beta = payload.beta if payload.kl else 0.0
|
| 61 |
+
policy_gradient = 0.0 if clipped else -advantage * ratio
|
| 62 |
+
kl_gradient = beta * (1 - reference_ratio)
|
| 63 |
+
logp_gradient = policy_gradient + kl_gradient
|
| 64 |
+
logit_gradient = weight * logp_gradient * (1 - current)
|
| 65 |
+
after = sigmoid(math.log(current) - math.log1p(-current) - payload.step_size * logit_gradient)
|
| 66 |
+
loss = -surrogate + beta * kl
|
| 67 |
+
return dict(old=old, reference=old, current=current, after=after,
|
| 68 |
+
delta=after-old, step_delta=after-current,
|
| 69 |
+
ratio=ratio, clipped_ratio=clipped_ratio, clipped=clipped,
|
| 70 |
+
kl=kl, policy_gradient=policy_gradient, kl_gradient=kl_gradient,
|
| 71 |
+
logp_gradient=logp_gradient, logit_gradient=logit_gradient,
|
| 72 |
+
weight=weight, surrogate=surrogate, loss=loss,
|
| 73 |
+
advantage=advantage, objective=-loss,
|
| 74 |
+
objective_contribution=-weight * loss,
|
| 75 |
+
objective_policy_gradient=-policy_gradient,
|
| 76 |
+
objective_kl_gradient=-kl_gradient,
|
| 77 |
+
objective_logp_gradient=-logp_gradient,
|
| 78 |
+
objective_weighted_logp_gradient=-weight * logp_gradient,
|
| 79 |
+
objective_logit_gradient=-logit_gradient)
|
| 80 |
+
|
| 81 |
+
|
| 82 |
def preview_step(rollouts, payload):
|
| 83 |
"""Differentiate the GRPO loss, then step each occurrence's sampled-token logit.
|
| 84 |
|
|
|
|
| 108 |
raise ValueError("A rollout has no sampled tokens.")
|
| 109 |
weight = 1 / (count * length)
|
| 110 |
for token in row["tokens"]:
|
| 111 |
+
detail = token_objective(token, advantage, weight, payload)
|
| 112 |
+
tokens.append(detail)
|
| 113 |
+
total_loss += weight * detail["loss"]
|
| 114 |
+
total_kl += weight * detail["kl"]
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 115 |
token_count += 1
|
| 116 |
+
clipped_count += int(detail["clipped"])
|
| 117 |
rows.append(dict(reward=reward, manual=override is not None, advantage=advantage,
|
| 118 |
scored_answer=scored_answer(row["text"]), tokens=tokens))
|
| 119 |
# J_GRPO is maximized; retain the equivalent loss L = -J for API compatibility.
|
| 120 |
return dict(mean=mean, std=std, objective=-total_loss, loss=total_loss, mean_kl=total_kl,
|
| 121 |
clipped_count=clipped_count, token_count=token_count,
|
| 122 |
all_equal=std == 0, weights_updated=False,
|
| 123 |
+
parameter_gradients_computed=False, updated_probabilities_computed=False,
|
| 124 |
probability_space="independent sampled-token logits", rollouts=rows)
|
| 125 |
|
| 126 |
|
| 127 |
class GRPOJobs:
|
| 128 |
+
algorithm = "grpo"
|
| 129 |
+
|
| 130 |
def __init__(self):
|
| 131 |
self.jobs = OrderedDict()
|
| 132 |
self.lock = threading.RLock()
|
|
|
|
| 156 |
raise BlockingIOError("This model is already computing. Try again when it finishes.")
|
| 157 |
sid = uuid.uuid4().hex
|
| 158 |
job = dict(id=sid, phase="sampling", message="Starting the rollout group…",
|
| 159 |
+
algorithm=self.algorithm,
|
| 160 |
context=context, prefix_tokens=[runtime.token(t) for t in ids],
|
| 161 |
config=payload.model_dump(), rollouts=[], elapsed=0.0,
|
| 162 |
weights_updated=False, **runtime.metadata(), _cancel=threading.Event())
|
|
|
|
| 164 |
with self.lock:
|
| 165 |
self.jobs[sid] = job
|
| 166 |
completed = [key for key, value in self.jobs.items()
|
| 167 |
+
if value["model"] == runtime.model_id and value["phase"] not in ("sampling", "valuing")]
|
| 168 |
for key in completed[:-2]:
|
| 169 |
del self.jobs[key]
|
| 170 |
threading.Thread(target=self._sample, args=(runtime, job, ids), daemon=True).start()
|
|
|
|
| 173 |
raise
|
| 174 |
return self.snapshot(sid)
|
| 175 |
|
| 176 |
+
def after_sampling(self, runtime, job, prefix_ids):
|
| 177 |
+
return not job["_cancel"].is_set()
|
| 178 |
+
|
| 179 |
+
def after_finished(self):
|
| 180 |
+
pass
|
| 181 |
+
|
| 182 |
@torch.inference_mode()
|
| 183 |
def _sample(self, runtime, job, prefix_ids):
|
| 184 |
started = time.perf_counter()
|
| 185 |
config = job["config"]
|
| 186 |
phase = "ready"
|
| 187 |
+
final_message = "Group complete. Assign rewards and inspect the objective gradients."
|
| 188 |
try:
|
| 189 |
generator = torch.Generator(device="cpu").manual_seed(config["seed"])
|
| 190 |
eos = runtime.model.generation_config.eos_token_id
|
|
|
|
| 202 |
final_message = "Sampling stopped. Model weights are unchanged."
|
| 203 |
return
|
| 204 |
message = f"Rollout {index+1} / {config['group_size']} · token {position+1} / {config['max_new_tokens']}"
|
| 205 |
+
runtime.active = f"{self.algorithm.upper()} · {message}"
|
| 206 |
with self.lock:
|
| 207 |
job.update(message=message, elapsed=round(time.perf_counter()-started, 2))
|
| 208 |
output = runtime.model(ids, past_key_values=cache, use_cache=True,
|
|
|
|
| 225 |
del cache
|
| 226 |
with self.lock:
|
| 227 |
row.update(finished=True, stop_reason="eos" if sampled[-1] in eos else "length")
|
| 228 |
+
if not self.after_sampling(runtime, job, prefix_ids):
|
| 229 |
+
phase = "cancelled"
|
| 230 |
+
final_message = "Computation stopped. Model weights are unchanged."
|
| 231 |
except Exception as exc:
|
| 232 |
import traceback
|
| 233 |
traceback.print_exc()
|
| 234 |
phase = "error"
|
| 235 |
+
final_message = f"Computation failed: {exc}"
|
| 236 |
finally:
|
| 237 |
with self.lock:
|
| 238 |
job.update(phase=phase, message=final_message, elapsed=round(time.perf_counter()-started, 2))
|
| 239 |
runtime.active = None
|
| 240 |
runtime.lock.release()
|
| 241 |
+
self.after_finished()
|
| 242 |
completed = [key for key, value in self.jobs.items()
|
| 243 |
+
if value["model"] == runtime.model_id and value["phase"] not in ("sampling", "valuing")]
|
| 244 |
for key in completed[:-2]:
|
| 245 |
del self.jobs[key]
|
ppo.py
ADDED
|
@@ -0,0 +1,126 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
"""Frozen-critic PPO: terminal rewards, per-token GAE, actor and critic losses."""
|
| 2 |
+
from __future__ import annotations
|
| 3 |
+
|
| 4 |
+
import math
|
| 5 |
+
import time
|
| 6 |
+
|
| 7 |
+
from grpo import GRPOJobs, automatic_reward, scored_answer, token_objective
|
| 8 |
+
|
| 9 |
+
|
| 10 |
+
def generalized_advantages(values, reward, gamma, gae_lambda):
|
| 11 |
+
"""Finite-horizon episodes: EOS AND the token budget terminate, V_T = 0."""
|
| 12 |
+
advantages, residuals = [0.0] * len(values), [0.0] * len(values)
|
| 13 |
+
running = 0.0
|
| 14 |
+
for t in reversed(range(len(values))):
|
| 15 |
+
next_value = values[t + 1] if t + 1 < len(values) else 0.0
|
| 16 |
+
delta = (reward if t == len(values)-1 else 0.0) + gamma * next_value - values[t]
|
| 17 |
+
running = delta + gamma * gae_lambda * running
|
| 18 |
+
residuals[t], advantages[t] = delta, running
|
| 19 |
+
return advantages, residuals
|
| 20 |
+
|
| 21 |
+
|
| 22 |
+
def ppo_objective(rollouts, payload):
|
| 23 |
+
count = len(rollouts)
|
| 24 |
+
if count < 2 or len(payload.rewards) != count:
|
| 25 |
+
raise ValueError("Provide one reward per rollout, with at least two rollouts.")
|
| 26 |
+
if payload.values is not None and len(payload.values) != count:
|
| 27 |
+
raise ValueError("Provide one value-override array per rollout.")
|
| 28 |
+
rewards, data = [], []
|
| 29 |
+
for index, (row, override) in enumerate(zip(rollouts, payload.rewards)):
|
| 30 |
+
reward = automatic_reward(row["text"], payload.answer, payload.reward_rule) if override is None else float(override)
|
| 31 |
+
values = list(row["values"])
|
| 32 |
+
if not values or len(values) != len(row["tokens"]):
|
| 33 |
+
raise ValueError("Each generated token needs one critic value.")
|
| 34 |
+
edits = payload.values[index] if payload.values is not None else None
|
| 35 |
+
if edits is not None:
|
| 36 |
+
if len(edits) != len(values):
|
| 37 |
+
raise ValueError("Value overrides must match the rollout's token count.")
|
| 38 |
+
values = [value if edit is None else float(edit) for value, edit in zip(values, edits)]
|
| 39 |
+
if not payload.critic:
|
| 40 |
+
values = [0.0] * len(values)
|
| 41 |
+
if not all(math.isfinite(x) for x in [reward, *values]):
|
| 42 |
+
raise ValueError("Rewards and critic values must be finite numbers.")
|
| 43 |
+
advantages, residuals = generalized_advantages(values, reward, payload.gamma, payload.gae_lambda)
|
| 44 |
+
rewards.append(reward)
|
| 45 |
+
data.append((values, advantages, residuals))
|
| 46 |
+
flat = [a for _, advantages, _ in data for a in advantages]
|
| 47 |
+
a_mean = sum(flat) / len(flat)
|
| 48 |
+
a_std = math.sqrt(sum((a-a_mean)**2 for a in flat) / max(1, len(flat)-1))
|
| 49 |
+
rows = []
|
| 50 |
+
objective = mean_kl = critic_loss = 0.0
|
| 51 |
+
clipped_count = 0
|
| 52 |
+
for row, reward, override, (values, advantages, residuals) in zip(rollouts, rewards, payload.rewards, data):
|
| 53 |
+
tokens = []
|
| 54 |
+
weight = 1 / (count * len(values))
|
| 55 |
+
for t, (token, value, raw_advantage, residual) in enumerate(zip(row["tokens"], values, advantages, residuals)):
|
| 56 |
+
advantage = (raw_advantage - a_mean) / (a_std + 1e-8) if payload.normalize else raw_advantage
|
| 57 |
+
detail = token_objective(token, advantage, weight, payload)
|
| 58 |
+
target = raw_advantage + value # Detached lambda-return target.
|
| 59 |
+
value_loss = 0.5 * (value - target)**2
|
| 60 |
+
detail.update(value=value, original_value=row["values"][t],
|
| 61 |
+
next_value=values[t+1] if t+1 < len(values) else 0.0,
|
| 62 |
+
reward=reward if t == len(values)-1 else 0.0,
|
| 63 |
+
terminal=t == len(values)-1, td_residual=residual,
|
| 64 |
+
raw_advantage=raw_advantage, return_target=target,
|
| 65 |
+
value_loss=value_loss, value_gradient=weight * (value-target))
|
| 66 |
+
tokens.append(detail)
|
| 67 |
+
objective += detail["objective_contribution"]
|
| 68 |
+
mean_kl += weight * detail["kl"]
|
| 69 |
+
critic_loss += weight * value_loss
|
| 70 |
+
clipped_count += int(detail["clipped"])
|
| 71 |
+
rows.append(dict(reward=reward, manual=override is not None,
|
| 72 |
+
scored_answer=scored_answer(row["text"]), tokens=tokens))
|
| 73 |
+
mean = sum(rewards)/count
|
| 74 |
+
std = math.sqrt(sum((r-mean)**2 for r in rewards)/(count-1))
|
| 75 |
+
return dict(mean=mean, std=std, objective=objective, loss=-objective,
|
| 76 |
+
critic_loss=critic_loss, mean_kl=mean_kl, clipped_count=clipped_count,
|
| 77 |
+
token_count=len(flat), all_equal=std == 0, advantage_mean=a_mean,
|
| 78 |
+
advantage_std=a_std, weights_updated=False, parameter_gradients_computed=False,
|
| 79 |
+
updated_probabilities_computed=False,
|
| 80 |
+
terminal_convention="EOS or token budget ends the episode; final reward at last action; V_T=0",
|
| 81 |
+
rollouts=rows)
|
| 82 |
+
|
| 83 |
+
|
| 84 |
+
class PPOJobs(GRPOJobs):
|
| 85 |
+
algorithm = "ppo"
|
| 86 |
+
|
| 87 |
+
def __init__(self, critic):
|
| 88 |
+
super().__init__()
|
| 89 |
+
self.critic = critic
|
| 90 |
+
|
| 91 |
+
def start(self, runtime, payload, context, ids):
|
| 92 |
+
if self.critic.phase != "ready":
|
| 93 |
+
raise ValueError(self.critic.message)
|
| 94 |
+
if not self.critic.lock.acquire(blocking=False):
|
| 95 |
+
raise BlockingIOError("The value critic is already computing. Try again when it finishes.")
|
| 96 |
+
try:
|
| 97 |
+
return super().start(runtime, payload, context, ids)
|
| 98 |
+
except Exception:
|
| 99 |
+
self.critic.lock.release()
|
| 100 |
+
raise
|
| 101 |
+
|
| 102 |
+
def after_sampling(self, runtime, job, prefix_ids):
|
| 103 |
+
with self.lock:
|
| 104 |
+
job.update(phase="valuing", critic=self.critic.status())
|
| 105 |
+
started = time.perf_counter() - job["elapsed"]
|
| 106 |
+
for index, row in enumerate(job["rollouts"]):
|
| 107 |
+
def progress(position):
|
| 108 |
+
message = f"Critic · rollout {index+1} / {len(job['rollouts'])} · state {position+1} / {len(row['tokens'])}"
|
| 109 |
+
runtime.active = f"PPO · {message}"
|
| 110 |
+
with self.lock:
|
| 111 |
+
job.update(message=message, elapsed=round(time.perf_counter()-started, 2))
|
| 112 |
+
values = self.critic.values(runtime, prefix_ids, row["tokens"], progress, job["_cancel"].is_set)
|
| 113 |
+
if values is None:
|
| 114 |
+
return False
|
| 115 |
+
with self.lock:
|
| 116 |
+
row["values"] = values
|
| 117 |
+
return not job["_cancel"].is_set()
|
| 118 |
+
|
| 119 |
+
def after_finished(self):
|
| 120 |
+
self.critic.lock.release()
|
| 121 |
+
|
| 122 |
+
def preview(self, sid, payload):
|
| 123 |
+
job = self.snapshot(sid)
|
| 124 |
+
if job["phase"] != "ready":
|
| 125 |
+
raise ValueError("Wait for the rollouts and critic values to finish.")
|
| 126 |
+
return ppo_objective(job["rollouts"], payload)
|
preload_models.py
CHANGED
|
@@ -2,10 +2,12 @@
|
|
| 2 |
from huggingface_hub import snapshot_download
|
| 3 |
|
| 4 |
from runtime import MODELS
|
|
|
|
| 5 |
|
| 6 |
|
| 7 |
if __name__ == "__main__":
|
| 8 |
-
|
|
|
|
| 9 |
snapshot_download(
|
| 10 |
repo_id=model_id,
|
| 11 |
revision=spec["revision"],
|
|
|
|
| 2 |
from huggingface_hub import snapshot_download
|
| 3 |
|
| 4 |
from runtime import MODELS
|
| 5 |
+
from critic import CRITIC_ID, CRITIC_REVISION
|
| 6 |
|
| 7 |
|
| 8 |
if __name__ == "__main__":
|
| 9 |
+
checkpoints = {**MODELS, CRITIC_ID: dict(revision=CRITIC_REVISION)}
|
| 10 |
+
for model_id, spec in checkpoints.items():
|
| 11 |
snapshot_download(
|
| 12 |
repo_id=model_id,
|
| 13 |
revision=spec["revision"],
|
server.py
CHANGED
|
@@ -15,11 +15,15 @@ from pydantic import BaseModel, ConfigDict, Field
|
|
| 15 |
|
| 16 |
from runtime import Runtime, MAX_TOKENS, MODEL_ID, GEMMA_ID, MODELS
|
| 17 |
from grpo import GRPOJobs
|
|
|
|
|
|
|
| 18 |
|
| 19 |
ROOT = Path(__file__).parent
|
| 20 |
ModelChoice = Literal[MODEL_ID, GEMMA_ID]
|
| 21 |
runtimes = {model_id: Runtime(model_id) for model_id in MODELS}
|
| 22 |
grpo_jobs = GRPOJobs()
|
|
|
|
|
|
|
| 23 |
|
| 24 |
|
| 25 |
class Message(BaseModel):
|
|
@@ -70,11 +74,21 @@ class PreviewInput(BaseModel):
|
|
| 70 |
drift: float = Field(default=0, ge=-3, le=3)
|
| 71 |
|
| 72 |
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 73 |
@asynccontextmanager
|
| 74 |
async def lifespan(app):
|
| 75 |
def load_models():
|
| 76 |
for runtime in runtimes.values():
|
| 77 |
runtime.load()
|
|
|
|
| 78 |
threading.Thread(target=load_models, daemon=True).start()
|
| 79 |
yield
|
| 80 |
|
|
@@ -82,21 +96,30 @@ async def lifespan(app):
|
|
| 82 |
app = FastAPI(lifespan=lifespan)
|
| 83 |
|
| 84 |
|
| 85 |
-
|
| 86 |
-
def start_grpo(payload: RolloutInput):
|
| 87 |
runtime = runtimes[payload.model]
|
| 88 |
try:
|
| 89 |
context, ids = runtime.prepare(Input(model=payload.model, mode=payload.mode, text=payload.prefix,
|
| 90 |
messages=[Message(role="user", content=payload.prefix)], thinking=payload.thinking))
|
| 91 |
if payload.mode == "chat" and payload.assistant_prefix:
|
| 92 |
context, ids = runtime.prepare(Input(model=payload.model, mode="raw", text=context+payload.assistant_prefix))
|
| 93 |
-
return
|
| 94 |
except BlockingIOError as exc:
|
| 95 |
raise HTTPException(409, str(exc)) from exc
|
| 96 |
except ValueError as exc:
|
| 97 |
raise HTTPException(422, str(exc)) from exc
|
| 98 |
|
| 99 |
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 100 |
@app.get("/api/grpo/runs/{sid}")
|
| 101 |
def get_grpo(sid: str):
|
| 102 |
try:
|
|
@@ -123,9 +146,35 @@ def preview_grpo(sid: str, payload: PreviewInput):
|
|
| 123 |
raise HTTPException(422, str(exc)) from exc
|
| 124 |
|
| 125 |
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 126 |
@app.get("/api/status")
|
| 127 |
def status(model: ModelChoice = MODEL_ID):
|
| 128 |
-
return dict(**runtimes[model].status(), models=[dict(id=mid, **spec) for mid, spec in MODELS.items()])
|
| 129 |
|
| 130 |
|
| 131 |
@app.post("/api/prepare")
|
|
@@ -235,6 +284,11 @@ def grpo_index():
|
|
| 235 |
return FileResponse(ROOT / "static" / "grpo.html")
|
| 236 |
|
| 237 |
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 238 |
app.mount("/", StaticFiles(directory=ROOT / "static"), name="static")
|
| 239 |
|
| 240 |
if __name__ == "__main__":
|
|
|
|
| 15 |
|
| 16 |
from runtime import Runtime, MAX_TOKENS, MODEL_ID, GEMMA_ID, MODELS
|
| 17 |
from grpo import GRPOJobs
|
| 18 |
+
from critic import CriticRuntime
|
| 19 |
+
from ppo import PPOJobs
|
| 20 |
|
| 21 |
ROOT = Path(__file__).parent
|
| 22 |
ModelChoice = Literal[MODEL_ID, GEMMA_ID]
|
| 23 |
runtimes = {model_id: Runtime(model_id) for model_id in MODELS}
|
| 24 |
grpo_jobs = GRPOJobs()
|
| 25 |
+
critic = CriticRuntime()
|
| 26 |
+
ppo_jobs = PPOJobs(critic)
|
| 27 |
|
| 28 |
|
| 29 |
class Message(BaseModel):
|
|
|
|
| 74 |
drift: float = Field(default=0, ge=-3, le=3)
|
| 75 |
|
| 76 |
|
| 77 |
+
class PPOPreviewInput(PreviewInput):
|
| 78 |
+
normalize: bool = False
|
| 79 |
+
step_size: float = Field(default=0, ge=0, le=100)
|
| 80 |
+
gamma: float = Field(default=1, ge=0, le=1)
|
| 81 |
+
gae_lambda: float = Field(default=0.95, ge=0, le=1)
|
| 82 |
+
critic: bool = True
|
| 83 |
+
values: list[Annotated[list[Annotated[float, Field(ge=-1000000, le=1000000)] | None], Field(max_length=64)] | None] | None = Field(default=None, max_length=12)
|
| 84 |
+
|
| 85 |
+
|
| 86 |
@asynccontextmanager
|
| 87 |
async def lifespan(app):
|
| 88 |
def load_models():
|
| 89 |
for runtime in runtimes.values():
|
| 90 |
runtime.load()
|
| 91 |
+
critic.load()
|
| 92 |
threading.Thread(target=load_models, daemon=True).start()
|
| 93 |
yield
|
| 94 |
|
|
|
|
| 96 |
app = FastAPI(lifespan=lifespan)
|
| 97 |
|
| 98 |
|
| 99 |
+
def start_rollouts(payload, jobs):
|
|
|
|
| 100 |
runtime = runtimes[payload.model]
|
| 101 |
try:
|
| 102 |
context, ids = runtime.prepare(Input(model=payload.model, mode=payload.mode, text=payload.prefix,
|
| 103 |
messages=[Message(role="user", content=payload.prefix)], thinking=payload.thinking))
|
| 104 |
if payload.mode == "chat" and payload.assistant_prefix:
|
| 105 |
context, ids = runtime.prepare(Input(model=payload.model, mode="raw", text=context+payload.assistant_prefix))
|
| 106 |
+
return jobs.start(runtime, payload, context, ids)
|
| 107 |
except BlockingIOError as exc:
|
| 108 |
raise HTTPException(409, str(exc)) from exc
|
| 109 |
except ValueError as exc:
|
| 110 |
raise HTTPException(422, str(exc)) from exc
|
| 111 |
|
| 112 |
|
| 113 |
+
@app.post("/api/grpo/runs")
|
| 114 |
+
def start_grpo(payload: RolloutInput):
|
| 115 |
+
return start_rollouts(payload, grpo_jobs)
|
| 116 |
+
|
| 117 |
+
|
| 118 |
+
@app.post("/api/ppo/runs")
|
| 119 |
+
def start_ppo(payload: RolloutInput):
|
| 120 |
+
return start_rollouts(payload, ppo_jobs)
|
| 121 |
+
|
| 122 |
+
|
| 123 |
@app.get("/api/grpo/runs/{sid}")
|
| 124 |
def get_grpo(sid: str):
|
| 125 |
try:
|
|
|
|
| 146 |
raise HTTPException(422, str(exc)) from exc
|
| 147 |
|
| 148 |
|
| 149 |
+
@app.get("/api/ppo/runs/{sid}")
|
| 150 |
+
def get_ppo(sid: str):
|
| 151 |
+
try:
|
| 152 |
+
return ppo_jobs.snapshot(sid)
|
| 153 |
+
except KeyError as exc:
|
| 154 |
+
raise HTTPException(404, "This PPO group has expired. Sample a new group.") from exc
|
| 155 |
+
|
| 156 |
+
|
| 157 |
+
@app.post("/api/ppo/runs/{sid}/cancel")
|
| 158 |
+
def cancel_ppo(sid: str):
|
| 159 |
+
try:
|
| 160 |
+
return ppo_jobs.cancel(sid)
|
| 161 |
+
except KeyError as exc:
|
| 162 |
+
raise HTTPException(404, "This PPO group has expired.") from exc
|
| 163 |
+
|
| 164 |
+
|
| 165 |
+
@app.post("/api/ppo/runs/{sid}/preview")
|
| 166 |
+
def preview_ppo(sid: str, payload: PPOPreviewInput):
|
| 167 |
+
try:
|
| 168 |
+
return ppo_jobs.preview(sid, payload)
|
| 169 |
+
except KeyError as exc:
|
| 170 |
+
raise HTTPException(404, "This PPO group has expired. Sample a new group.") from exc
|
| 171 |
+
except ValueError as exc:
|
| 172 |
+
raise HTTPException(422, str(exc)) from exc
|
| 173 |
+
|
| 174 |
+
|
| 175 |
@app.get("/api/status")
|
| 176 |
def status(model: ModelChoice = MODEL_ID):
|
| 177 |
+
return dict(**runtimes[model].status(), models=[dict(id=mid, **spec) for mid, spec in MODELS.items()], critic=critic.status())
|
| 178 |
|
| 179 |
|
| 180 |
@app.post("/api/prepare")
|
|
|
|
| 284 |
return FileResponse(ROOT / "static" / "grpo.html")
|
| 285 |
|
| 286 |
|
| 287 |
+
@app.get("/ppo")
|
| 288 |
+
def ppo_index():
|
| 289 |
+
return FileResponse(ROOT / "static" / "ppo.html")
|
| 290 |
+
|
| 291 |
+
|
| 292 |
app.mount("/", StaticFiles(directory=ROOT / "static"), name="static")
|
| 293 |
|
| 294 |
if __name__ == "__main__":
|
static/grpo.css
CHANGED
|
@@ -35,7 +35,8 @@
|
|
| 35 |
.g-ablation{padding:10px;border:1px solid var(--line);background:#fff;border-radius:5px;cursor:pointer;min-width:0}
|
| 36 |
.g-ablation>span{display:flex;align-items:center;gap:8px;font-size:12px;font-weight:600}
|
| 37 |
.g-ablation small{display:block;color:var(--muted);font-size:11px;line-height:1.5;margin:6px 0 0 23px}
|
| 38 |
-
.g-
|
|
|
|
| 39 |
.g-slider-grid label{display:block;font-size:11px;color:var(--muted);min-width:0}
|
| 40 |
.g-slider-grid output{float:right;color:var(--ink);font:12px var(--mono);margin-left:4px}
|
| 41 |
.g-slider-grid input{display:block;width:100%;margin:12px 0 0;accent-color:var(--accent);height:16px}
|
|
@@ -73,9 +74,15 @@
|
|
| 73 |
.g-token:hover{border-color:var(--accent)}.g-token.active{outline:2px solid var(--accent);outline-offset:1px}
|
| 74 |
.g-token-token{font:600 12px var(--mono);max-width:100%;overflow:hidden;text-overflow:ellipsis;white-space:nowrap}
|
| 75 |
.g-token-prob{font:10px var(--mono);color:#555;white-space:nowrap}
|
| 76 |
-
.g-token-
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 77 |
.g-token.up{background:#fff;border-color:#aaa}.g-token.down{background:#fff;border-color:#aaa}
|
| 78 |
-
.g-token.up .g-token-
|
| 79 |
.g-token.is-clipped{border-style:dashed;border-color:#555}
|
| 80 |
.g-completion{margin-top:12px;color:var(--muted);font-size:11px}.g-completion summary{cursor:pointer}.g-completion pre{white-space:pre-wrap;overflow-wrap:anywhere;max-height:180px;overflow:auto;padding:10px;background:var(--soft);font:12px/1.7 var(--mono)}
|
| 81 |
.g-prefix-details{margin:18px 0;color:var(--muted);font-size:12px}.g-prefix-details summary{cursor:pointer}.g-prefix-details pre{white-space:pre-wrap;overflow-wrap:anywhere;font:12px/1.7 var(--mono);max-height:240px;overflow:auto;padding:12px;background:var(--soft);border-radius:4px}
|
|
@@ -91,6 +98,7 @@
|
|
| 91 |
.g-equation-scroll{min-width:0;max-width:100%;overflow-x:auto;padding:12px 4px;scrollbar-width:thin}
|
| 92 |
.g-equation-scroll math{font-size:21px;width:max-content;margin:0 auto}
|
| 93 |
@media(max-width:600px){.g-equation-scroll math{font-size:18px}}
|
|
|
|
| 94 |
.grpo-page .site-footer{gap:15px;flex-wrap:wrap}
|
| 95 |
@media(min-width:1500px){.g-workspace{grid-template-columns:310px minmax(0,1fr);gap:38px}}
|
| 96 |
@media(max-width:1100px){.g-workspace{grid-template-columns:265px minmax(0,1fr);gap:22px}.g-slider-grid{grid-template-columns:repeat(2,minmax(0,1fr))}.g-ablation-grid{grid-template-columns:1fr}.g-ablation{display:flex;justify-content:space-between;align-items:center;gap:12px}.g-ablation small{margin:0;text-align:right}.g-token-metrics{grid-template-columns:repeat(2,minmax(0,1fr))}.g-process{display:none}}
|
|
|
|
| 35 |
.g-ablation{padding:10px;border:1px solid var(--line);background:#fff;border-radius:5px;cursor:pointer;min-width:0}
|
| 36 |
.g-ablation>span{display:flex;align-items:center;gap:8px;font-size:12px;font-weight:600}
|
| 37 |
.g-ablation small{display:block;color:var(--muted);font-size:11px;line-height:1.5;margin:6px 0 0 23px}
|
| 38 |
+
.g-ppo-ablations{grid-template-columns:repeat(2,minmax(0,1fr))}
|
| 39 |
+
.g-slider-grid{display:grid;grid-template-columns:repeat(3,minmax(0,1fr));gap:20px;margin-top:19px}
|
| 40 |
.g-slider-grid label{display:block;font-size:11px;color:var(--muted);min-width:0}
|
| 41 |
.g-slider-grid output{float:right;color:var(--ink);font:12px var(--mono);margin-left:4px}
|
| 42 |
.g-slider-grid input{display:block;width:100%;margin:12px 0 0;accent-color:var(--accent);height:16px}
|
|
|
|
| 74 |
.g-token:hover{border-color:var(--accent)}.g-token.active{outline:2px solid var(--accent);outline-offset:1px}
|
| 75 |
.g-token-token{font:600 12px var(--mono);max-width:100%;overflow:hidden;text-overflow:ellipsis;white-space:nowrap}
|
| 76 |
.g-token-prob{font:10px var(--mono);color:#555;white-space:nowrap}
|
| 77 |
+
.g-token-gradient{font:600 11px var(--mono);white-space:nowrap}
|
| 78 |
+
.g-token-value{font:10px var(--mono);color:#555;white-space:nowrap}
|
| 79 |
+
.g-value-editor{display:flex;align-items:end;gap:14px;flex-wrap:wrap;margin:15px 0 8px}
|
| 80 |
+
.g-value-editor label{font-size:12px;color:var(--muted)}
|
| 81 |
+
.g-value-editor input[type=number]{max-width:180px}
|
| 82 |
+
.g-critic-note{border-left:2px solid #999;padding-left:12px}
|
| 83 |
+
.g-ppo-summary{grid-template-columns:repeat(4,minmax(0,1fr))}
|
| 84 |
.g-token.up{background:#fff;border-color:#aaa}.g-token.down{background:#fff;border-color:#aaa}
|
| 85 |
+
.g-token.up .g-token-gradient{color:#222}.g-token.down .g-token-gradient{color:#222}
|
| 86 |
.g-token.is-clipped{border-style:dashed;border-color:#555}
|
| 87 |
.g-completion{margin-top:12px;color:var(--muted);font-size:11px}.g-completion summary{cursor:pointer}.g-completion pre{white-space:pre-wrap;overflow-wrap:anywhere;max-height:180px;overflow:auto;padding:10px;background:var(--soft);font:12px/1.7 var(--mono)}
|
| 88 |
.g-prefix-details{margin:18px 0;color:var(--muted);font-size:12px}.g-prefix-details summary{cursor:pointer}.g-prefix-details pre{white-space:pre-wrap;overflow-wrap:anywhere;font:12px/1.7 var(--mono);max-height:240px;overflow:auto;padding:12px;background:var(--soft);border-radius:4px}
|
|
|
|
| 98 |
.g-equation-scroll{min-width:0;max-width:100%;overflow-x:auto;padding:12px 4px;scrollbar-width:thin}
|
| 99 |
.g-equation-scroll math{font-size:21px;width:max-content;margin:0 auto}
|
| 100 |
@media(max-width:600px){.g-equation-scroll math{font-size:18px}}
|
| 101 |
+
@media(max-width:820px){.g-ppo-summary{grid-template-columns:repeat(2,minmax(0,1fr));gap:18px}.g-ppo-summary article:nth-child(3){border:0;padding-left:0}}
|
| 102 |
.grpo-page .site-footer{gap:15px;flex-wrap:wrap}
|
| 103 |
@media(min-width:1500px){.g-workspace{grid-template-columns:310px minmax(0,1fr);gap:38px}}
|
| 104 |
@media(max-width:1100px){.g-workspace{grid-template-columns:265px minmax(0,1fr);gap:22px}.g-slider-grid{grid-template-columns:repeat(2,minmax(0,1fr))}.g-ablation-grid{grid-template-columns:1fr}.g-ablation{display:flex;justify-content:space-between;align-items:center;gap:12px}.g-ablation small{margin:0;text-align:right}.g-token-metrics{grid-template-columns:repeat(2,minmax(0,1fr))}.g-process{display:none}}
|
static/grpo.html
CHANGED
|
@@ -13,12 +13,12 @@
|
|
| 13 |
<body class="grpo-page">
|
| 14 |
<header class="site-header">
|
| 15 |
<a class="brand" href="/"><span class="brand-mark">L</span> Layer lab <span class="brand-light">/</span><span class="brand-caption">GRPO sandbox</span></a>
|
| 16 |
-
<nav><a class="mode-link" href="/">Layers & attention</a><a class="mode-link active" href="/grpo" aria-current="page">GRPO</a><span class="model-pill">Weights frozen</span></nav>
|
| 17 |
</header>
|
| 18 |
<main>
|
| 19 |
<div class="intro">
|
| 20 |
-
<div><p class="eyebrow">GROUP RELATIVE POLICY OPTIMIZATION</p><h1>One group.
|
| 21 |
-
<div class="g-process" aria-label="Experiment steps"><span><b>01</b> Sample</span><i>→</i><span><b>02</b> Reward</span><i>→</i><span><b>03</b>
|
| 22 |
</div>
|
| 23 |
<div class="g-workspace">
|
| 24 |
<aside class="controls g-controls" aria-label="Rollout settings">
|
|
@@ -64,7 +64,7 @@
|
|
| 64 |
</aside>
|
| 65 |
<div class="g-results">
|
| 66 |
<section class="g-step-panel" aria-labelledby="g-step-title">
|
| 67 |
-
<div class="g-section-heading"><div><p class="eyebrow">03 / THE OBJECTIVE</p><h2 id="g-step-title">
|
| 68 |
<div class="g-ablation-grid">
|
| 69 |
<label class="g-ablation"><span><input id="g-normalize" type="checkbox" checked> Scale advantages Âᵢ</span><small>(rᵢ − mean(r)) / (std(r) + δ)</small></label>
|
| 70 |
<label class="g-ablation"><span><input id="g-clipping" type="checkbox" checked> Clip the ratio ρᵢ,ₜ</span><small>ρᶜˡⁱᵖᵢ,ₜ = clip(ρᵢ,ₜ, 1−ε, 1+ε)</small></label>
|
|
@@ -73,11 +73,10 @@
|
|
| 73 |
<div class="g-slider-grid">
|
| 74 |
<label>Clip ε <output id="g-epsilon-value">0.20</output><input id="g-epsilon" type="range" min="0.01" max="0.5" step="0.01" value="0.2"></label>
|
| 75 |
<label>KL β <output id="g-beta-value">0.04</output><input id="g-beta" type="range" min="0" max="1" step="0.01" value="0.04"></label>
|
| 76 |
-
<label>Preview step size η <output id="g-step-value">10.0</output><input id="g-step" type="range" min="0" max="100" step="0.5" value="10"></label>
|
| 77 |
<label>Candidate drift d <output id="g-drift-value">0.00</output><input id="g-drift" type="range" min="-3" max="3" step="0.05" value="0"></label>
|
| 78 |
</div>
|
| 79 |
<p id="g-origin-note" class="g-explanation">At d = 0, πθ = πold = πref: ρᵢ,ₜ = 1 and Kᵢ,ₜ = 0. Clipping and KL do not change this first gradient. Move drift in either direction to explore their effects.</p>
|
| 80 |
-
<
|
| 81 |
</section>
|
| 82 |
<div class="g-summary" aria-live="polite">
|
| 83 |
<article><span>Mean reward · mean(r)</span><strong id="g-mean">—</strong><small id="g-std">Group baseline</small></article>
|
|
@@ -89,13 +88,13 @@
|
|
| 89 |
<div id="g-progress" class="g-progress hidden" role="status"><span class="spinner"></span><span id="g-progress-text">Sampling…</span></div>
|
| 90 |
<p id="g-group-note" class="g-explanation">Choose a prefix and target, then sample a group. No completions are prewritten.</p>
|
| 91 |
<p id="g-equal" class="g-warning hidden">All rewards rᵢ are equal → every advantage Âᵢ is zero. Edit a reward or resample to see a learning signal.</p>
|
| 92 |
-
<div class="g-token-legend"><span id="g-token-legend">Each token shows its collection probability.</span><span><i class="g-swatch" aria-hidden="true">+</i>
|
| 93 |
<div id="g-rollouts" class="g-rollouts"><div class="g-empty"><span>G</span><h3>A group makes the comparison.</h3><p>Five real rollouts by default.<br>Rewards become advantages relative to this group.</p></div></div>
|
| 94 |
<details id="g-prefix-details" class="g-prefix-details hidden"><summary>Inspect the exact collection prefix</summary><pre id="g-formatted-prefix"></pre><div id="g-prefix-tokens" class="g-prefix-tokens"></div></details>
|
| 95 |
</section>
|
| 96 |
<section id="g-inspector" class="g-inspector hidden" aria-labelledby="g-inspector-title"><div class="g-section-heading"><div><p class="eyebrow">SELECTED TOKEN</p><h2 id="g-inspector-title">Follow the gradient</h2></div></div><div id="g-token-detail"></div></section>
|
| 97 |
<section class="g-method" id="g-method">
|
| 98 |
-
<p class="eyebrow">WHAT THIS
|
| 99 |
<p>Rollouts and collection probabilities come from the selected checkpoint. The reference stays equal to that frozen collection policy. Rewards are assigned to complete rollouts; their group-relative advantages apply to every generated token, including the stop token. The prefix receives no reward gradient.</p>
|
| 100 |
<div class="g-equations" aria-label="GRPO objective and notation">
|
| 101 |
<div class="g-equation-scroll"><math xmlns="http://www.w3.org/1998/Math/MathML" display="block" aria-label="J GRPO is the expectation of the group and token averaged minimum of rho times A hat and clipped rho times A hat, minus beta K."><mrow><msub><mi>𝒥</mi><mtext>GRPO</mtext></msub><mo>=</mo><mi>𝔼</mi><mo>[</mo><mfrac><mi>1</mi><mi>G</mi></mfrac><munder><mo>∑</mo><mi>i</mi></munder><mfrac><mi>1</mi><mrow><mo>|</mo><msub><mi>y</mi><mi>i</mi></msub><mo>|</mo></mrow></mfrac><munder><mo>∑</mo><mi>t</mi></munder><mrow><mo>(</mo><mrow><mrow><mi mathvariant="normal">min</mi><mrow><mo>(</mo><mrow><msub><mi>ρ</mi><mrow><mi>i</mi><mo>,</mo><mi>t</mi></mrow></msub><msub><mover accent="true"><mi>A</mi><mo>^</mo></mover><mi>i</mi></msub><mo>,</mo><msubsup><mi>ρ</mi><mrow><mi>i</mi><mo>,</mo><mi>t</mi></mrow><mtext>clip</mtext></msubsup><msub><mover accent="true"><mi>A</mi><mo>^</mo></mover><mi>i</mi></msub></mrow><mo>)</mo></mrow></mrow><mo>−</mo><mi>β</mi><msub><mi>K</mi><mrow><mi>i</mi><mo>,</mo><mi>t</mi></mrow></msub></mrow><mo>)</mo></mrow><mo>]</mo></mrow></math></div>
|
|
@@ -104,10 +103,11 @@
|
|
| 104 |
<div class="g-equation-scroll"><math xmlns="http://www.w3.org/1998/Math/MathML" display="block" aria-label="A hat i is reward r i minus the group mean reward, divided by the reward standard deviation plus delta."><mrow><msub><mover accent="true"><mi>A</mi><mo>^</mo></mover><mi>i</mi></msub><mo>=</mo><mfrac><mrow><msub><mi>r</mi><mi>i</mi></msub><mo>−</mo><mrow><mi mathvariant="normal">mean</mi><mrow><mo>(</mo><mi mathvariant="bold">r</mi><mo>)</mo></mrow></mrow></mrow><mrow><mrow><mi mathvariant="normal">std</mi><mrow><mo>(</mo><mi mathvariant="bold">r</mi><mo>)</mo></mrow></mrow><mo>+</mo><mi>δ</mi></mrow></mfrac></mrow></math></div>
|
| 105 |
</div>
|
| 106 |
<p>Here G is the number of rollouts, yᵢ is a generated completion, and hᵢ,ₜ is the prefix plus its earlier tokens. Reward rᵢ belongs to the whole rollout. πθ denotes the candidate policy and πold the collection policy. The frozen reference πref equals πold. Kᵢ,ₜ is the sampled KL estimate <code>πref/πθ − log(πref/πθ) − 1</code> for the same token and history. We use sample standard deviation and δ = 10⁻⁸; disabling scaling leaves <code>Âᵢ = rᵢ − mean(r)</code>. The displayed 𝒥GRPO is one group's estimate of the expectation.</p>
|
| 107 |
-
<p>
|
| 108 |
-
<p>
|
|
|
|
| 109 |
</section>
|
| 110 |
-
<footer class="site-footer"><span>Layer lab <i>/</i> GRPO sandbox</span><span>
|
| 111 |
</div>
|
| 112 |
</div>
|
| 113 |
</main>
|
|
|
|
| 13 |
<body class="grpo-page">
|
| 14 |
<header class="site-header">
|
| 15 |
<a class="brand" href="/"><span class="brand-mark">L</span> Layer lab <span class="brand-light">/</span><span class="brand-caption">GRPO sandbox</span></a>
|
| 16 |
+
<nav><a class="mode-link" href="/">Layers & attention</a><a class="mode-link active" href="/grpo" aria-current="page">GRPO</a><a class="mode-link" href="/ppo">PPO</a><span class="model-pill">Weights frozen</span></nav>
|
| 17 |
</header>
|
| 18 |
<main>
|
| 19 |
<div class="intro">
|
| 20 |
+
<div><p class="eyebrow">GROUP RELATIVE POLICY OPTIMIZATION</p><h1>One group. Follow the gradient.</h1><p class="intro-description">Sample real completions, assign rewards, and inspect each token's learning signal.<br>Exact objective derivatives. Frozen model weights.</p></div>
|
| 21 |
+
<div class="g-process" aria-label="Experiment steps"><span><b>01</b> Sample</span><i>→</i><span><b>02</b> Reward</span><i>→</i><span><b>03</b> Gradient</span></div>
|
| 22 |
</div>
|
| 23 |
<div class="g-workspace">
|
| 24 |
<aside class="controls g-controls" aria-label="Rollout settings">
|
|
|
|
| 64 |
</aside>
|
| 65 |
<div class="g-results">
|
| 66 |
<section class="g-step-panel" aria-labelledby="g-step-title">
|
| 67 |
+
<div class="g-section-heading"><div><p class="eyebrow">03 / THE OBJECTIVE</p><h2 id="g-step-title">Inspect the objective</h2></div><button id="g-export" class="text-button" disabled>Export JSON ↓</button></div>
|
| 68 |
<div class="g-ablation-grid">
|
| 69 |
<label class="g-ablation"><span><input id="g-normalize" type="checkbox" checked> Scale advantages Âᵢ</span><small>(rᵢ − mean(r)) / (std(r) + δ)</small></label>
|
| 70 |
<label class="g-ablation"><span><input id="g-clipping" type="checkbox" checked> Clip the ratio ρᵢ,ₜ</span><small>ρᶜˡⁱᵖᵢ,ₜ = clip(ρᵢ,ₜ, 1−ε, 1+ε)</small></label>
|
|
|
|
| 73 |
<div class="g-slider-grid">
|
| 74 |
<label>Clip ε <output id="g-epsilon-value">0.20</output><input id="g-epsilon" type="range" min="0.01" max="0.5" step="0.01" value="0.2"></label>
|
| 75 |
<label>KL β <output id="g-beta-value">0.04</output><input id="g-beta" type="range" min="0" max="1" step="0.01" value="0.04"></label>
|
|
|
|
| 76 |
<label>Candidate drift d <output id="g-drift-value">0.00</output><input id="g-drift" type="range" min="-3" max="3" step="0.05" value="0"></label>
|
| 77 |
</div>
|
| 78 |
<p id="g-origin-note" class="g-explanation">At d = 0, πθ = πold = πref: ρᵢ,ₜ = 1 and Kᵢ,ₜ = 0. Clipping and KL do not change this first gradient. Move drift in either direction to explore their effects.</p>
|
| 79 |
+
<p id="g-preview-status" class="field-note" role="status">Sample a group to see gradients on every token automatically.</p>
|
| 80 |
</section>
|
| 81 |
<div class="g-summary" aria-live="polite">
|
| 82 |
<article><span>Mean reward · mean(r)</span><strong id="g-mean">—</strong><small id="g-std">Group baseline</small></article>
|
|
|
|
| 88 |
<div id="g-progress" class="g-progress hidden" role="status"><span class="spinner"></span><span id="g-progress-text">Sampling…</span></div>
|
| 89 |
<p id="g-group-note" class="g-explanation">Choose a prefix and target, then sample a group. No completions are prewritten.</p>
|
| 90 |
<p id="g-equal" class="g-warning hidden">All rewards rᵢ are equal → every advantage Âᵢ is zero. Edit a reward or resample to see a learning signal.</p>
|
| 91 |
+
<div class="g-token-legend"><span id="g-token-legend">Each token shows its collection probability and objective gradient.</span><span><i class="g-swatch" aria-hidden="true">+</i> Positive g <i class="g-swatch negative" aria-hidden="true">−</i> Negative g <i class="g-swatch clipped"></i> Clipped</span></div>
|
| 92 |
<div id="g-rollouts" class="g-rollouts"><div class="g-empty"><span>G</span><h3>A group makes the comparison.</h3><p>Five real rollouts by default.<br>Rewards become advantages relative to this group.</p></div></div>
|
| 93 |
<details id="g-prefix-details" class="g-prefix-details hidden"><summary>Inspect the exact collection prefix</summary><pre id="g-formatted-prefix"></pre><div id="g-prefix-tokens" class="g-prefix-tokens"></div></details>
|
| 94 |
</section>
|
| 95 |
<section id="g-inspector" class="g-inspector hidden" aria-labelledby="g-inspector-title"><div class="g-section-heading"><div><p class="eyebrow">SELECTED TOKEN</p><h2 id="g-inspector-title">Follow the gradient</h2></div></div><div id="g-token-detail"></div></section>
|
| 96 |
<section class="g-method" id="g-method">
|
| 97 |
+
<p class="eyebrow">WHAT THIS SANDBOX COMPUTES</p><h2>Real samples. Exact objective derivatives.</h2>
|
| 98 |
<p>Rollouts and collection probabilities come from the selected checkpoint. The reference stays equal to that frozen collection policy. Rewards are assigned to complete rollouts; their group-relative advantages apply to every generated token, including the stop token. The prefix receives no reward gradient.</p>
|
| 99 |
<div class="g-equations" aria-label="GRPO objective and notation">
|
| 100 |
<div class="g-equation-scroll"><math xmlns="http://www.w3.org/1998/Math/MathML" display="block" aria-label="J GRPO is the expectation of the group and token averaged minimum of rho times A hat and clipped rho times A hat, minus beta K."><mrow><msub><mi>𝒥</mi><mtext>GRPO</mtext></msub><mo>=</mo><mi>𝔼</mi><mo>[</mo><mfrac><mi>1</mi><mi>G</mi></mfrac><munder><mo>∑</mo><mi>i</mi></munder><mfrac><mi>1</mi><mrow><mo>|</mo><msub><mi>y</mi><mi>i</mi></msub><mo>|</mo></mrow></mfrac><munder><mo>∑</mo><mi>t</mi></munder><mrow><mo>(</mo><mrow><mrow><mi mathvariant="normal">min</mi><mrow><mo>(</mo><mrow><msub><mi>ρ</mi><mrow><mi>i</mi><mo>,</mo><mi>t</mi></mrow></msub><msub><mover accent="true"><mi>A</mi><mo>^</mo></mover><mi>i</mi></msub><mo>,</mo><msubsup><mi>ρ</mi><mrow><mi>i</mi><mo>,</mo><mi>t</mi></mrow><mtext>clip</mtext></msubsup><msub><mover accent="true"><mi>A</mi><mo>^</mo></mover><mi>i</mi></msub></mrow><mo>)</mo></mrow></mrow><mo>−</mo><mi>β</mi><msub><mi>K</mi><mrow><mi>i</mi><mo>,</mo><mi>t</mi></mrow></msub></mrow><mo>)</mo></mrow><mo>]</mo></mrow></math></div>
|
|
|
|
| 103 |
<div class="g-equation-scroll"><math xmlns="http://www.w3.org/1998/Math/MathML" display="block" aria-label="A hat i is reward r i minus the group mean reward, divided by the reward standard deviation plus delta."><mrow><msub><mover accent="true"><mi>A</mi><mo>^</mo></mover><mi>i</mi></msub><mo>=</mo><mfrac><mrow><msub><mi>r</mi><mi>i</mi></msub><mo>−</mo><mrow><mi mathvariant="normal">mean</mi><mrow><mo>(</mo><mi mathvariant="bold">r</mi><mo>)</mo></mrow></mrow></mrow><mrow><mrow><mi mathvariant="normal">std</mi><mrow><mo>(</mo><mi mathvariant="bold">r</mi><mo>)</mo></mrow></mrow><mo>+</mo><mi>δ</mi></mrow></mfrac></mrow></math></div>
|
| 104 |
</div>
|
| 105 |
<p>Here G is the number of rollouts, yᵢ is a generated completion, and hᵢ,ₜ is the prefix plus its earlier tokens. Reward rᵢ belongs to the whole rollout. πθ denotes the candidate policy and πold the collection policy. The frozen reference πref equals πold. Kᵢ,ₜ is the sampled KL estimate <code>πref/πθ − log(πref/πθ) − 1</code> for the same token and history. We use sample standard deviation and δ = 10⁻⁸; disabling scaling leaves <code>Âᵢ = rᵢ − mean(r)</code>. The displayed 𝒥GRPO is one group's estimate of the expectation.</p>
|
| 106 |
+
<p>Each chip shows the collection probability and <code>gᵢ,ₜ = ∂𝒥GRPO/∂log πθ(yᵢ,ₜ | hᵢ,ₜ)</code>, including group/token averaging. Positive g locally favors increasing this log-probability; negative g favors decreasing it. Selecting a token also shows its objective contribution and equivalent loss. The loss is the negative of the maximized objective.</p>
|
| 107 |
+
<p>Candidate drift sets <code>logit(πθ) = logit(πold) + dÂᵢ</code>, with d from −3 to +3. This is a synthetic candidate for exploring clipping and KL, not a model forward pass or a measured training change. Positive drift moves with the advantage; negative drift moves against it. At zero drift the candidate is exactly the collection policy.</p>
|
| 108 |
+
<p><strong>A gradient is not a probability change.</strong> Actual training backpropagates <code>∇θ𝒥 = Σ gᵢ,ₜ ∇θ log πθ</code>, updates shared weights with an optimizer, then computes new probabilities with another forward pass. Shared parameters couple tokens, and optimizer state affects the result. This sandbox computes no parameter gradients or optimizer steps and reports no updated-model probabilities. GRPO clipping removes an incentive; it does not forcibly clamp probabilities.</p>
|
| 109 |
</section>
|
| 110 |
+
<footer class="site-footer"><span>Layer lab <i>/</i> GRPO sandbox</span><span>Samples · rewards · objective gradients · frozen weights</span></footer>
|
| 111 |
</div>
|
| 112 |
</div>
|
| 113 |
</main>
|
static/grpo.js
CHANGED
|
@@ -4,7 +4,11 @@ const visible = text => String(text).replace(/ /g,'·').replace(/\n/g,'↵').rep
|
|
| 4 |
const pct = p => `${(p*100).toFixed(p>=.1?1:p>=.001?2:3)}%`;
|
| 5 |
const number = n => Number(n).toFixed(4);
|
| 6 |
const signed = (n, digits=2) => `${n<0?'−':'+'}${Math.abs(n).toFixed(digits)}`;
|
| 7 |
-
const
|
|
|
|
|
|
|
|
|
|
|
|
|
| 8 |
let statusTimer, jobTimer, previewTimer, previewVersion=0, editor=null;
|
| 9 |
|
| 10 |
async function api(path, body) {
|
|
@@ -23,34 +27,37 @@ function collectionPayload(){
|
|
| 23 |
function objectivePayload(){
|
| 24 |
return {answer:$('#g-answer').value, reward_rule:$('#g-rule').value, rewards:state.rewards,
|
| 25 |
normalize:$('#g-normalize').checked, clipping:$('#g-clipping').checked, epsilon:Number($('#g-epsilon').value),
|
| 26 |
-
kl:$('#g-kl').checked, beta:Number($('#g-beta').value), step_size:
|
|
|
|
| 27 |
}
|
| 28 |
function updateControls(){
|
| 29 |
-
const sampling=state.busy||state.job
|
| 30 |
document.querySelectorAll('[data-sampling]').forEach(el=>{el.disabled=sampling;});
|
| 31 |
$('#g-example').disabled=sampling;
|
| 32 |
const raw=$('#g-mode').value==='raw', native=$('#g-model').value==='Qwen/Qwen3-0.6B';
|
| 33 |
$('#g-assistant-field').classList.toggle('hidden',raw);
|
| 34 |
$('#g-thinking').disabled=sampling||raw||!native;
|
| 35 |
$('#g-thinking-note').textContent=!native?'Gemma learns Reasoning / Solution text; no native thinking switch.':raw?'Reasoning is controlled by the exact Raw prefix.':"Qwen's native thinking template.";
|
| 36 |
-
$('#g-sample').disabled=sampling||state.status?.phase!=='ready'||Boolean(state.status?.active);
|
| 37 |
$('#g-sample span:first-child').textContent=sampling?'Sampling…':`Sample ${$('#g-count').value||5} rollouts`;
|
| 38 |
-
$('#g-cancel').classList.toggle('hidden',state.job
|
| 39 |
-
$('#g-
|
| 40 |
-
$('#g-preview span:first-child').textContent=state.showStep?'Preview shown':'Preview one step';
|
| 41 |
-
$('#g-hide-preview').classList.toggle('hidden',!state.showStep);
|
| 42 |
$('#g-reset-rewards').disabled=sampling||!ready;
|
| 43 |
$('#g-export').disabled=sampling||!ready||!state.preview||state.previewBusy;
|
| 44 |
document.querySelectorAll('[data-reward]').forEach(el=>{el.disabled=sampling||!ready;});
|
| 45 |
$('#g-epsilon').disabled=!$('#g-clipping').checked;
|
| 46 |
$('#g-beta').disabled=!$('#g-kl').checked;
|
| 47 |
-
$('#g-token-legend').textContent=
|
| 48 |
-
for(const [id,digits] of [['epsilon',2],['beta',2],['
|
| 49 |
const drift=Number($('#g-drift').value);
|
| 50 |
$('#g-drift-value').textContent=drift===0?'0.00':signed(drift);
|
| 51 |
$('#g-origin-note').textContent=Number($('#g-drift').value)===0?
|
| 52 |
'At d = 0, πθ = πold = πref: ρᵢ,ₜ = 1 and Kᵢ,ₜ = 0. Clipping and KL do not change this first gradient. Move drift in either direction to explore their effects.':
|
| 53 |
-
|
|
|
|
|
|
|
|
|
|
|
|
|
| 54 |
}
|
| 55 |
function samplingChanged(){
|
| 56 |
if(state.job)$('#g-stale').classList.remove('hidden');
|
|
@@ -70,23 +77,23 @@ async function pollStatus(){
|
|
| 70 |
$('#g-status').textContent=result.active||result.message;
|
| 71 |
renderSpecial(result);updateControls();
|
| 72 |
}catch(exc){if(model===$('#g-model').value){state.status=null;$('#g-status').textContent=`Cannot reach the server: ${exc.message}`;updateControls();}}
|
| 73 |
-
if(model===$('#g-model').value)statusTimer=setTimeout(pollStatus,state.job?.phase==
|
| 74 |
}
|
| 75 |
|
| 76 |
function tokenMarkup(token, ri, ti){
|
| 77 |
-
const detail=state.preview?.rollouts[ri]?.tokens[ti], show=
|
| 78 |
-
const change=
|
| 79 |
const direction=change>1e-10?'up':change< -1e-10?'down':'';
|
| 80 |
-
const title=`Token ${ti} · ID ${token.id} · BPE ${token.piece}\nCollected πold: ${pct(token.probability)}${detail?`\
|
| 81 |
-
return `<button type="button" class="g-token ${show?direction:''} ${show&&detail.clipped?'is-clipped':''} ${state.selected?.[0]===ri&&state.selected?.[1]===ti?'active':''}" data-token="${ri}:${ti}" data-old="${token.probability}" ${detail?`data-
|
| 82 |
}
|
| 83 |
function renderCards(){
|
| 84 |
const job=state.job;if(!job)return;
|
| 85 |
-
$('#g-rollouts').innerHTML=job.rollouts.map((row,index)=>`<article class="g-rollout" data-rollout="${index}"><div class="g-rollout-head"><div class="g-rollout-id">Rollout ${String(index+1).padStart(2,'0')} <small>${row.tokens.length} tokens · ${row.finished?(row.stop_reason==='eos'?'stop token':'token limit'):'sampling…'}</small>${row.stop_reason==='length'?'<small class="g-truncated">Truncated</small>':''}</div><div class="g-reward-fields"><label><span>Reward r<sub>${index+1}</sub></span> <input aria-label="Reward for rollout ${index+1}" type="number" step="any" min="-1000000" max="1000000" value="0" data-reward="${index}" ${job.phase==='ready'?'':'disabled'}></label><span class="g-advantage" data-advantage="${index}" title="Group-relative advantage Âᵢ">Âᵢ —</span></div></div><div class="g-rollout-body"><p class="g-answer-line" data-answer="${index}">Reward assigned after the group completes.</p><div class="g-token-strip" data-strip="${index}">${row.tokens.map((token,ti)=>tokenMarkup(token,index,ti)).join('')}</div><details class="g-completion"><summary>Completion text</summary><pre>${escape(row.text||'(No visible text yet)')}</pre></details></div></article>`).join('');
|
| 86 |
if(state.preview)renderPreview();
|
| 87 |
}
|
| 88 |
function renderJob(){
|
| 89 |
-
const job=state.job, sampling=job
|
| 90 |
$('#g-progress').classList.toggle('hidden',!sampling);
|
| 91 |
$('#g-progress-text').textContent=`${job.message} · ${job.elapsed.toFixed(1)} s`;
|
| 92 |
$('#g-group-note').textContent=`${job.name} · G = ${job.config.group_size} rollouts · seed ${job.config.seed} · temperature ${job.config.temperature} · ${job.elapsed.toFixed(1)} s. ${sampling?'Sampling from the fixed checkpoint.':job.phase==='ready'?'Rewards can be changed without sampling again.':job.message}`;
|
|
@@ -102,16 +109,18 @@ function renderPreview(){
|
|
| 102 |
$('#g-loss').textContent=number(preview.objective);
|
| 103 |
$('#g-clipped-count').textContent=`${preview.clipped_count} / ${preview.token_count}`;
|
| 104 |
$('#g-kl-metric').textContent=`Mean Kᵢ,ₜ: ${number(preview.mean_kl)}${$('#g-kl').checked?'':' · penalty off'}`;
|
| 105 |
-
$('#g-equal').classList.toggle('hidden',!preview.all_equal);
|
| 106 |
preview.rollouts.forEach((row,index)=>{
|
| 107 |
const input=$(`[data-reward="${index}"]`);
|
| 108 |
if(input&&document.activeElement!==input)input.value=row.reward;
|
| 109 |
const adv=$(`[data-advantage="${index}"]`);
|
| 110 |
-
adv.innerHTML=`Â<sub>${index+1}</sub> ${signed(row.advantage,3)}`;adv.className=`g-advantage ${row.advantage>0?'g-positive':row.advantage<0?'g-negative':''}`;
|
|
|
|
| 111 |
$(`[data-answer="${index}"]`).innerHTML=`Scored answer: <strong>${escape(row.scored_answer||'(empty / unfinished reasoning)')}</strong>${row.manual?'<span class="g-manual-tag">MANUAL REWARD</span>':` · ${$('#g-rule').value==='exact'?'exact match':'contains answer'}`}`;
|
| 112 |
$(`[data-strip="${index}"]`).innerHTML=state.job.rollouts[index].tokens.map((token,ti)=>tokenMarkup(token,index,ti)).join('');
|
| 113 |
});
|
| 114 |
-
$('#g-preview-status').textContent=
|
|
|
|
| 115 |
renderInspector();updateControls();
|
| 116 |
}
|
| 117 |
function renderInspector(){
|
|
@@ -120,8 +129,10 @@ function renderInspector(){
|
|
| 120 |
if(!detail)return;
|
| 121 |
const token=state.job.rollouts[selected[0]].tokens[selected[1]], row=state.preview.rollouts[selected[0]];
|
| 122 |
$('#g-inspector-title').textContent=`Rollout ${selected[0]+1} · token ${selected[1]} · ${visible(token.text)}`;
|
| 123 |
-
const metrics=[['Collection πold',pct(detail.old)],['
|
| 124 |
-
|
|
|
|
|
|
|
| 125 |
}
|
| 126 |
|
| 127 |
function schedulePreview(){
|
|
@@ -132,37 +143,37 @@ function schedulePreview(){
|
|
| 132 |
previewTimer=setTimeout(async()=>{
|
| 133 |
try{
|
| 134 |
if(!$('#g-answer').value.trim())throw new Error('Enter a correct answer to calculate automatic rewards.');
|
| 135 |
-
const result=await api(`/
|
| 136 |
if(version!==previewVersion||state.job?.id!==sid)return;
|
| 137 |
state.preview=result;error();renderPreview();
|
| 138 |
-
}catch(exc){if(version===previewVersion&&state.job?.id===sid){state.preview=null;error(exc.message);$('#g-preview-status').textContent='
|
| 139 |
finally{if(version===previewVersion){state.previewBusy=false;$('#g-rollouts').setAttribute('aria-busy','false');updateControls();}}
|
| 140 |
},180);
|
| 141 |
}
|
| 142 |
async function pollJob(sid){
|
| 143 |
clearTimeout(jobTimer);
|
| 144 |
try{
|
| 145 |
-
const job=await api(`/
|
| 146 |
if(state.job?.id!==sid)return;
|
| 147 |
state.job=job;renderJob();
|
| 148 |
-
if(job
|
| 149 |
else{
|
| 150 |
state.status={...state.status,active:null};updateControls();
|
| 151 |
-
if(job.phase==='ready'){state.rewards=job.rollouts.map(()=>null);schedulePreview();}
|
| 152 |
else if(job.phase==='error')error(job.message);
|
| 153 |
}
|
| 154 |
}catch(exc){if(state.job?.id===sid){error(`Could not check the rollout group: ${exc.message}`);jobTimer=setTimeout(()=>pollJob(sid),2500);}}
|
| 155 |
}
|
| 156 |
$('#g-form').addEventListener('submit',async event=>{
|
| 157 |
-
event.preventDefault();if(state.busy||state.job
|
| 158 |
if(!$('#g-form').reportValidity())return;
|
| 159 |
const payload=collectionPayload();
|
| 160 |
error();state.busy=true;updateControls();
|
| 161 |
try{
|
| 162 |
-
const job=await api(
|
| 163 |
clearTimeout(jobTimer);clearTimeout(previewTimer);previewVersion++;
|
| 164 |
-
state.job=job;state.preview=null;state.rewards=[];state.
|
| 165 |
-
for(const id of ['g-mean','g-loss','g-clipped-count'])$('#'+id).textContent='—';
|
| 166 |
$('#g-std').textContent='Group baseline';$('#g-kl-metric').textContent='Mean Kᵢ,ₜ: —';
|
| 167 |
$('#g-stale').classList.add('hidden');$('#g-equal').classList.add('hidden');$('#g-inspector').classList.add('hidden');
|
| 168 |
$('#g-preview-status').textContent='Waiting for the rollout group…';renderJob();pollJob(job.id);
|
|
@@ -171,11 +182,9 @@ $('#g-form').addEventListener('submit',async event=>{
|
|
| 171 |
});
|
| 172 |
$('#g-cancel').addEventListener('click',async()=>{
|
| 173 |
if(!state.job)return;$('#g-cancel').disabled=true;
|
| 174 |
-
try{await api(`/
|
| 175 |
finally{$('#g-cancel').disabled=false;}
|
| 176 |
});
|
| 177 |
-
$('#g-preview').addEventListener('click',()=>{state.showStep=true;renderPreview();});
|
| 178 |
-
$('#g-hide-preview').addEventListener('click',()=>{state.showStep=false;renderPreview();});
|
| 179 |
$('#g-reset-rewards').addEventListener('click',()=>{state.rewards=state.job.rollouts.map(()=>null);schedulePreview();});
|
| 180 |
$('#g-rollouts').addEventListener('input',event=>{
|
| 181 |
const input=event.target.closest('[data-reward]');if(!input)return;
|
|
@@ -188,27 +197,38 @@ $('#g-rollouts').addEventListener('click',event=>{
|
|
| 188 |
state.selected=token.dataset.token.split(':').map(Number);
|
| 189 |
document.querySelectorAll('.g-token.active').forEach(el=>el.classList.remove('active'));token.classList.add('active');renderInspector();
|
| 190 |
});
|
| 191 |
-
for(const id of ['g-answer','g-rule','g-normalize','g-clipping','g-kl','g-epsilon','g-beta','g-
|
| 192 |
document.querySelectorAll('[data-sampling]').forEach(el=>el.addEventListener('input',samplingChanged));
|
| 193 |
$('#g-model').addEventListener('change',()=>{state.status=null;$('#g-thinking').checked=false;updateControls();pollStatus();});
|
| 194 |
for(const id of ['g-prefix','g-assistant'])$('#'+id).addEventListener('focus',event=>{editor=event.target;});
|
| 195 |
$('#g-special-tokens').addEventListener('click',event=>{
|
| 196 |
-
const button=event.target.closest('[data-insert]');if(!button||state.job
|
| 197 |
const target=$('#g-mode').value==='raw'||!editor||editor.disabled?$('#g-prefix'):editor;
|
| 198 |
target.focus();target.setRangeText(button.dataset.insert,target.selectionStart,target.selectionEnd,'end');target.dispatchEvent(new Event('input',{bubbles:true}));
|
| 199 |
});
|
| 200 |
$('#g-example').addEventListener('click',()=>{
|
| 201 |
-
$('#g-mode').value='chat';$('#g-prefix').value='Name one primary color. Reply with one word.';$('#g-assistant').value='';$('#g-answer').value='red';$('#g-rule').value='exact';
|
| 202 |
$('#g-count').value='5';$('#g-length').value='24';$('#g-temperature').value='1';$('#g-seed').value='42';$('#g-thinking').checked=false;
|
| 203 |
samplingChanged();schedulePreview();
|
| 204 |
});
|
| 205 |
$('#g-export').addEventListener('click',()=>{
|
| 206 |
if(!state.job||!state.preview)return;
|
| 207 |
-
const data={kind:
|
| 208 |
-
notation:{objective:
|
| 209 |
-
note:'Real frozen-policy samples
|
| 210 |
const url=URL.createObjectURL(new Blob([JSON.stringify(data,null,2)],{type:'application/json'}));
|
| 211 |
-
const a=document.createElement('a');a.href=url;a.download=`
|
| 212 |
});
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 213 |
document.addEventListener('keydown',event=>{if((event.metaKey||event.ctrlKey)&&event.key==='Enter'){event.preventDefault();if(!$('#g-sample').disabled)$('#g-form').requestSubmit();}});
|
| 214 |
updateControls();pollStatus();
|
|
|
|
| 4 |
const pct = p => `${(p*100).toFixed(p>=.1?1:p>=.001?2:3)}%`;
|
| 5 |
const number = n => Number(n).toFixed(4);
|
| 6 |
const signed = (n, digits=2) => `${n<0?'−':'+'}${Math.abs(n).toFixed(digits)}`;
|
| 7 |
+
const signal = n => n===0?'0':`${n<0?'−':'+'}${Math.abs(n).toPrecision(3)}`;
|
| 8 |
+
const PPO = location.pathname==='/ppo';
|
| 9 |
+
const algorithm = PPO?'PPO':'GRPO', apiRoot=`/api/${algorithm.toLowerCase()}/runs`;
|
| 10 |
+
const computing = job => ['sampling','valuing'].includes(job?.phase);
|
| 11 |
+
const state = {job:null, preview:null, rewards:[], values:[], selected:null, status:null, busy:false, previewBusy:false};
|
| 12 |
let statusTimer, jobTimer, previewTimer, previewVersion=0, editor=null;
|
| 13 |
|
| 14 |
async function api(path, body) {
|
|
|
|
| 27 |
function objectivePayload(){
|
| 28 |
return {answer:$('#g-answer').value, reward_rule:$('#g-rule').value, rewards:state.rewards,
|
| 29 |
normalize:$('#g-normalize').checked, clipping:$('#g-clipping').checked, epsilon:Number($('#g-epsilon').value),
|
| 30 |
+
kl:$('#g-kl').checked, beta:Number($('#g-beta').value), step_size:0, drift:Number($('#g-drift').value),
|
| 31 |
+
...(PPO?{gamma:Number($('#g-gamma').value), gae_lambda:Number($('#g-lambda').value), critic:$('#g-critic').checked, values:state.values}: {})};
|
| 32 |
}
|
| 33 |
function updateControls(){
|
| 34 |
+
const sampling=state.busy||computing(state.job), ready=state.job?.phase==='ready';
|
| 35 |
document.querySelectorAll('[data-sampling]').forEach(el=>{el.disabled=sampling;});
|
| 36 |
$('#g-example').disabled=sampling;
|
| 37 |
const raw=$('#g-mode').value==='raw', native=$('#g-model').value==='Qwen/Qwen3-0.6B';
|
| 38 |
$('#g-assistant-field').classList.toggle('hidden',raw);
|
| 39 |
$('#g-thinking').disabled=sampling||raw||!native;
|
| 40 |
$('#g-thinking-note').textContent=!native?'Gemma learns Reasoning / Solution text; no native thinking switch.':raw?'Reasoning is controlled by the exact Raw prefix.':"Qwen's native thinking template.";
|
| 41 |
+
$('#g-sample').disabled=sampling||state.status?.phase!=='ready'||Boolean(state.status?.active)||(PPO&&(state.status?.critic?.phase!=='ready'||state.status?.critic?.busy));
|
| 42 |
$('#g-sample span:first-child').textContent=sampling?'Sampling…':`Sample ${$('#g-count').value||5} rollouts`;
|
| 43 |
+
$('#g-cancel').classList.toggle('hidden',!computing(state.job));
|
| 44 |
+
$('#g-cancel').textContent=state.job?.phase==='valuing'?'Stop critic evaluation':'Stop sampling';
|
|
|
|
|
|
|
| 45 |
$('#g-reset-rewards').disabled=sampling||!ready;
|
| 46 |
$('#g-export').disabled=sampling||!ready||!state.preview||state.previewBusy;
|
| 47 |
document.querySelectorAll('[data-reward]').forEach(el=>{el.disabled=sampling||!ready;});
|
| 48 |
$('#g-epsilon').disabled=!$('#g-clipping').checked;
|
| 49 |
$('#g-beta').disabled=!$('#g-kl').checked;
|
| 50 |
+
$('#g-token-legend').textContent='πold = collected probability · g = ∂𝒥/∂log πθ, including averaging · g is not a probability change';
|
| 51 |
+
for(const [id,digits] of [['epsilon',2],['beta',2],...(PPO?[['gamma',2],['lambda',2]]:[])])$('#g-'+id+'-value').textContent=Number($('#g-'+id).value).toFixed(digits);
|
| 52 |
const drift=Number($('#g-drift').value);
|
| 53 |
$('#g-drift-value').textContent=drift===0?'0.00':signed(drift);
|
| 54 |
$('#g-origin-note').textContent=Number($('#g-drift').value)===0?
|
| 55 |
'At d = 0, πθ = πold = πref: ρᵢ,ₜ = 1 and Kᵢ,ₜ = 0. Clipping and KL do not change this first gradient. Move drift in either direction to explore their effects.':
|
| 56 |
+
`Illustrative candidate drift shifts each sampled-token logit by d × ${PPO?'Âᵢ,ₜ':'Âᵢ'}. Positive d moves with the advantage; negative d moves against it. This synthetic candidate helps explore clipping and KL; it is not a model update.`;
|
| 57 |
+
if(PPO){
|
| 58 |
+
$('#g-critic-status').textContent=state.status?.critic?.message||'Connecting to the value critic…';
|
| 59 |
+
$('#g-reset-values').disabled=sampling||!ready;
|
| 60 |
+
}
|
| 61 |
}
|
| 62 |
function samplingChanged(){
|
| 63 |
if(state.job)$('#g-stale').classList.remove('hidden');
|
|
|
|
| 77 |
$('#g-status').textContent=result.active||result.message;
|
| 78 |
renderSpecial(result);updateControls();
|
| 79 |
}catch(exc){if(model===$('#g-model').value){state.status=null;$('#g-status').textContent=`Cannot reach the server: ${exc.message}`;updateControls();}}
|
| 80 |
+
if(model===$('#g-model').value)statusTimer=setTimeout(pollStatus,computing(state.job)||state.status?.phase!=='ready'||(PPO&&state.status?.critic?.phase!=='ready')?1800:5000);
|
| 81 |
}
|
| 82 |
|
| 83 |
function tokenMarkup(token, ri, ti){
|
| 84 |
+
const detail=state.preview?.rollouts[ri]?.tokens[ti], show=Boolean(detail);
|
| 85 |
+
const change=detail?.objective_weighted_logp_gradient||0;
|
| 86 |
const direction=change>1e-10?'up':change< -1e-10?'down':'';
|
| 87 |
+
const title=`Token ${ti} · ID ${token.id} · BPE ${token.piece}\nCollected πold: ${pct(token.probability)}${detail?`\nGradient g = ∂𝒥/∂log πθ: ${change}\nToken objective (before averaging): ${detail.objective}\nIllustrative candidate πθ: ${pct(detail.current)}\nRatio ρᵢ,ₜ: ${number(detail.ratio)} · ${detail.clipped?'policy gradient clipped':'policy gradient active'}${PPO?`\nV(h): ${detail.value}\nGAE advantage: ${detail.advantage}\nTD residual: ${detail.td_residual}`:''}\nNo probability update or model-weight gradient is computed.`:''}`;
|
| 88 |
+
return `<button type="button" class="g-token ${show?direction:''} ${show&&detail.clipped?'is-clipped':''} ${state.selected?.[0]===ri&&state.selected?.[1]===ti?'active':''}" data-token="${ri}:${ti}" data-old="${token.probability}" ${detail?`data-gradient="${change}" data-clipped="${detail.clipped}"`:''} title="${escape(title)}" aria-label="${escape(title)}">${show?`<span class="g-token-gradient">g ${signal(change)}</span>`:''}<span class="g-token-token">${escape(visible(token.text))}</span><span class="g-token-prob">πold ${pct(token.probability)}</span>${PPO&&show?`<span class="g-token-value">V ${number(detail.value)}</span><span class="g-token-value">Â ${signal(detail.advantage)}</span>`:''}</button>`;
|
| 89 |
}
|
| 90 |
function renderCards(){
|
| 91 |
const job=state.job;if(!job)return;
|
| 92 |
+
$('#g-rollouts').innerHTML=job.rollouts.map((row,index)=>`<article class="g-rollout" data-rollout="${index}"><div class="g-rollout-head"><div class="g-rollout-id">Rollout ${String(index+1).padStart(2,'0')} <small>${row.tokens.length} tokens · ${row.finished?(row.stop_reason==='eos'?'stop token':'token limit'):'sampling…'}</small>${row.stop_reason==='length'?'<small class="g-truncated">Truncated</small>':''}</div><div class="g-reward-fields"><label><span>Reward r<sub>${index+1}</sub></span> <input aria-label="Reward for rollout ${index+1}" type="number" step="any" min="-1000000" max="1000000" value="0" data-reward="${index}" ${job.phase==='ready'?'':'disabled'}></label><span class="g-advantage" data-advantage="${index}" title="${PPO?'Per-token state values and GAE advantages':'Group-relative advantage Âᵢ'}">${PPO?'V / Â pending':'Âᵢ —'}</span></div></div><div class="g-rollout-body"><p class="g-answer-line" data-answer="${index}">Reward assigned after the group completes.</p><div class="g-token-strip" data-strip="${index}">${row.tokens.map((token,ti)=>tokenMarkup(token,index,ti)).join('')}</div><details class="g-completion"><summary>Completion text</summary><pre>${escape(row.text||'(No visible text yet)')}</pre></details></div></article>`).join('');
|
| 93 |
if(state.preview)renderPreview();
|
| 94 |
}
|
| 95 |
function renderJob(){
|
| 96 |
+
const job=state.job, sampling=computing(job);
|
| 97 |
$('#g-progress').classList.toggle('hidden',!sampling);
|
| 98 |
$('#g-progress-text').textContent=`${job.message} · ${job.elapsed.toFixed(1)} s`;
|
| 99 |
$('#g-group-note').textContent=`${job.name} · G = ${job.config.group_size} rollouts · seed ${job.config.seed} · temperature ${job.config.temperature} · ${job.elapsed.toFixed(1)} s. ${sampling?'Sampling from the fixed checkpoint.':job.phase==='ready'?'Rewards can be changed without sampling again.':job.message}`;
|
|
|
|
| 109 |
$('#g-loss').textContent=number(preview.objective);
|
| 110 |
$('#g-clipped-count').textContent=`${preview.clipped_count} / ${preview.token_count}`;
|
| 111 |
$('#g-kl-metric').textContent=`Mean Kᵢ,ₜ: ${number(preview.mean_kl)}${$('#g-kl').checked?'':' · penalty off'}`;
|
| 112 |
+
$('#g-equal').classList.toggle('hidden',PPO||!preview.all_equal);
|
| 113 |
preview.rollouts.forEach((row,index)=>{
|
| 114 |
const input=$(`[data-reward="${index}"]`);
|
| 115 |
if(input&&document.activeElement!==input)input.value=row.reward;
|
| 116 |
const adv=$(`[data-advantage="${index}"]`);
|
| 117 |
+
adv.innerHTML=PPO?'Per-token V & Â below':`Â<sub>${index+1}</sub> ${signed(row.advantage,3)}`;adv.className=`g-advantage ${row.advantage>0?'g-positive':row.advantage<0?'g-negative':''}`;
|
| 118 |
+
if(PPO)adv.title='PPO uses a separate state value and GAE advantage at every token.';
|
| 119 |
$(`[data-answer="${index}"]`).innerHTML=`Scored answer: <strong>${escape(row.scored_answer||'(empty / unfinished reasoning)')}</strong>${row.manual?'<span class="g-manual-tag">MANUAL REWARD</span>':` · ${$('#g-rule').value==='exact'?'exact match':'contains answer'}`}`;
|
| 120 |
$(`[data-strip="${index}"]`).innerHTML=state.job.rollouts[index].tokens.map((token,ti)=>tokenMarkup(token,index,ti)).join('');
|
| 121 |
});
|
| 122 |
+
$('#g-preview-status').textContent='Objective and token gradients update automatically. Probabilities and weights stay fixed.';
|
| 123 |
+
if(PPO)$('#g-critic-loss').textContent=number(preview.critic_loss);
|
| 124 |
renderInspector();updateControls();
|
| 125 |
}
|
| 126 |
function renderInspector(){
|
|
|
|
| 129 |
if(!detail)return;
|
| 130 |
const token=state.job.rollouts[selected[0]].tokens[selected[1]], row=state.preview.rollouts[selected[0]];
|
| 131 |
$('#g-inspector-title').textContent=`Rollout ${selected[0]+1} · token ${selected[1]} · ${visible(token.text)}`;
|
| 132 |
+
const metrics=[['Collection πold',pct(detail.old)],['Illustrative candidate πθ',pct(detail.current)],['Token objective jᵢ,ₜ',number(detail.objective)],['Token loss −jᵢ,ₜ',number(detail.loss)],[PPO?'GAE advantage Âᵢ,ₜ':'Advantage Âᵢ',signed(detail.advantage,4)],['Policy ratio ρᵢ,ₜ',number(detail.ratio)],['Sampled KL Kᵢ,ₜ',number(detail.kl)],['Policy branch',detail.clipped?'Clipped · zero slope':'Unclipped'],['Gradient g = ∂𝒥/∂log πθ',signal(detail.objective_weighted_logp_gradient)],...(PPO?[['Value V(hᵢ,ₜ)',number(detail.value)],['Next value V(hᵢ,ₜ₊₁)',number(detail.next_value)],['Immediate reward',number(detail.reward)],['TD residual δᵢ,ₜ',number(detail.td_residual)],['Raw GAE advantage',number(detail.raw_advantage)],['λ-return target',number(detail.return_target)],['Value loss ½(V − R̂)²',number(detail.value_loss)]]:[])];
|
| 133 |
+
const editing=document.activeElement?.id==='g-value-edit';
|
| 134 |
+
$('#g-token-detail').innerHTML=`<div class="g-token-metrics">${metrics.map(([label,value])=>`<div><span>${label}</span><strong>${escape(value)}</strong></div>`).join('')}</div><p class="g-gradient-line">Before averaging: policy contribution to <code>∂j/∂log πθ = ${number(detail.objective_policy_gradient)}</code> · KL contribution <code>${number(detail.objective_kl_gradient)}</code>.<br>The chip shows <code>g = ∂𝒥${algorithm}/∂log πθ = ${signal(detail.objective_weighted_logp_gradient)}</code>, including averaging weight <code>${detail.weight.toPrecision(4)}</code>.<br>A positive g locally favors increasing this log-probability; a negative g favors decreasing it. Actual training uses <code>∇θ𝒥 = Σ gᵢ,ₜ ∇θ log πθ</code>, an optimizer step, and a new forward pass. Shared model weights couple all tokens, so g does not predict Δ probability.${detail.clipped?' The clipped policy term has zero slope here; an enabled KL term can still contribute.':''}</p>${PPO?`<div class="g-value-editor"><label for="g-value-edit">Override this state's V <input id="g-value-edit" type="number" step="any" min="-1000000" max="1000000" value="${detail.value}" ${$('#g-critic').checked?'':'disabled'}></label><button type="button" id="g-reset-value" class="text-button">Use checkpoint value (${number(detail.original_value)})</button></div><p class="field-note">Value is evaluated before this token. ${detail.terminal?'This action ends the finite-horizon episode, so next V = 0.':'Next V comes from the history after this token.'} Critic loss uses a detached λ-return; neither network is updated.</p>`:''}`;
|
| 135 |
+
if(editing)$('#g-value-edit')?.focus();
|
| 136 |
}
|
| 137 |
|
| 138 |
function schedulePreview(){
|
|
|
|
| 143 |
previewTimer=setTimeout(async()=>{
|
| 144 |
try{
|
| 145 |
if(!$('#g-answer').value.trim())throw new Error('Enter a correct answer to calculate automatic rewards.');
|
| 146 |
+
const result=await api(`${apiRoot}/${sid}/preview`,objectivePayload());
|
| 147 |
if(version!==previewVersion||state.job?.id!==sid)return;
|
| 148 |
state.preview=result;error();renderPreview();
|
| 149 |
+
}catch(exc){if(version===previewVersion&&state.job?.id===sid){state.preview=null;renderCards();renderInspector();error(exc.message);$('#g-preview-status').textContent='Objective unavailable. Check the reward settings.';}}
|
| 150 |
finally{if(version===previewVersion){state.previewBusy=false;$('#g-rollouts').setAttribute('aria-busy','false');updateControls();}}
|
| 151 |
},180);
|
| 152 |
}
|
| 153 |
async function pollJob(sid){
|
| 154 |
clearTimeout(jobTimer);
|
| 155 |
try{
|
| 156 |
+
const job=await api(`${apiRoot}/${sid}`);
|
| 157 |
if(state.job?.id!==sid)return;
|
| 158 |
state.job=job;renderJob();
|
| 159 |
+
if(computing(job))jobTimer=setTimeout(()=>pollJob(sid),800);
|
| 160 |
else{
|
| 161 |
state.status={...state.status,active:null};updateControls();
|
| 162 |
+
if(job.phase==='ready'){state.rewards=job.rollouts.map(()=>null);state.values=job.rollouts.map(row=>row.tokens.map(()=>null));schedulePreview();}
|
| 163 |
else if(job.phase==='error')error(job.message);
|
| 164 |
}
|
| 165 |
}catch(exc){if(state.job?.id===sid){error(`Could not check the rollout group: ${exc.message}`);jobTimer=setTimeout(()=>pollJob(sid),2500);}}
|
| 166 |
}
|
| 167 |
$('#g-form').addEventListener('submit',async event=>{
|
| 168 |
+
event.preventDefault();if(state.busy||computing(state.job))return;
|
| 169 |
if(!$('#g-form').reportValidity())return;
|
| 170 |
const payload=collectionPayload();
|
| 171 |
error();state.busy=true;updateControls();
|
| 172 |
try{
|
| 173 |
+
const job=await api(apiRoot,payload);
|
| 174 |
clearTimeout(jobTimer);clearTimeout(previewTimer);previewVersion++;
|
| 175 |
+
state.job=job;state.preview=null;state.rewards=[];state.values=[];state.selected=null;state.previewBusy=false;
|
| 176 |
+
for(const id of ['g-mean','g-loss','g-clipped-count',...(PPO?['g-critic-loss']:[])])$('#'+id).textContent='—';
|
| 177 |
$('#g-std').textContent='Group baseline';$('#g-kl-metric').textContent='Mean Kᵢ,ₜ: —';
|
| 178 |
$('#g-stale').classList.add('hidden');$('#g-equal').classList.add('hidden');$('#g-inspector').classList.add('hidden');
|
| 179 |
$('#g-preview-status').textContent='Waiting for the rollout group…';renderJob();pollJob(job.id);
|
|
|
|
| 182 |
});
|
| 183 |
$('#g-cancel').addEventListener('click',async()=>{
|
| 184 |
if(!state.job)return;$('#g-cancel').disabled=true;
|
| 185 |
+
try{await api(`${apiRoot}/${state.job.id}/cancel`,{});}catch(exc){error(exc.message);}
|
| 186 |
finally{$('#g-cancel').disabled=false;}
|
| 187 |
});
|
|
|
|
|
|
|
| 188 |
$('#g-reset-rewards').addEventListener('click',()=>{state.rewards=state.job.rollouts.map(()=>null);schedulePreview();});
|
| 189 |
$('#g-rollouts').addEventListener('input',event=>{
|
| 190 |
const input=event.target.closest('[data-reward]');if(!input)return;
|
|
|
|
| 197 |
state.selected=token.dataset.token.split(':').map(Number);
|
| 198 |
document.querySelectorAll('.g-token.active').forEach(el=>el.classList.remove('active'));token.classList.add('active');renderInspector();
|
| 199 |
});
|
| 200 |
+
for(const id of ['g-answer','g-rule','g-normalize','g-clipping','g-kl','g-epsilon','g-beta','g-drift',...(PPO?['g-gamma','g-lambda','g-critic']:[])])$('#'+id).addEventListener('input',()=>{updateControls();schedulePreview();});
|
| 201 |
document.querySelectorAll('[data-sampling]').forEach(el=>el.addEventListener('input',samplingChanged));
|
| 202 |
$('#g-model').addEventListener('change',()=>{state.status=null;$('#g-thinking').checked=false;updateControls();pollStatus();});
|
| 203 |
for(const id of ['g-prefix','g-assistant'])$('#'+id).addEventListener('focus',event=>{editor=event.target;});
|
| 204 |
$('#g-special-tokens').addEventListener('click',event=>{
|
| 205 |
+
const button=event.target.closest('[data-insert]');if(!button||computing(state.job))return;
|
| 206 |
const target=$('#g-mode').value==='raw'||!editor||editor.disabled?$('#g-prefix'):editor;
|
| 207 |
target.focus();target.setRangeText(button.dataset.insert,target.selectionStart,target.selectionEnd,'end');target.dispatchEvent(new Event('input',{bubbles:true}));
|
| 208 |
});
|
| 209 |
$('#g-example').addEventListener('click',()=>{
|
| 210 |
+
$('#g-mode').value='chat';$('#g-prefix').value=PPO?'What is 2 + 3? Reply with only the answer.':'Name one primary color. Reply with one word.';$('#g-assistant').value='';$('#g-answer').value=PPO?'5':'red';$('#g-rule').value='exact';
|
| 211 |
$('#g-count').value='5';$('#g-length').value='24';$('#g-temperature').value='1';$('#g-seed').value='42';$('#g-thinking').checked=false;
|
| 212 |
samplingChanged();schedulePreview();
|
| 213 |
});
|
| 214 |
$('#g-export').addEventListener('click',()=>{
|
| 215 |
if(!state.job||!state.preview)return;
|
| 216 |
+
const data={kind:`${algorithm.toLowerCase()}-objective-gradients`,weights_updated:false,collection:state.job,objective:objectivePayload(),preview:state.preview,
|
| 217 |
+
notation:{objective:`J_${algorithm} (maximized)`,reward:'r_i',advantage:PPO?'A_hat_i_t (GAE)':'A_hat_i',gradient:'g_i_t = d J / d log pi_theta (including averaging)',ratio:'rho_i_t',clipped_ratio:'rho_i_t^clip',kl:'K_i_t',current:'illustrative pi_theta(y_i_t | h_i_t)',old:'pi_old(y_i_t | h_i_t)',reference:'pi_ref(y_i_t | h_i_t)',normalization_epsilon:'delta = 1e-8',loss:`L = -J_${algorithm} (minimized)`},
|
| 218 |
+
note:'Real frozen-policy samples and objective derivatives. No model parameter gradients, optimizer updates, or updated-model probabilities. Legacy after/delta fields describe an independent-logit toy calculation only; the UI displays g, never a predicted probability change.'};
|
| 219 |
const url=URL.createObjectURL(new Blob([JSON.stringify(data,null,2)],{type:'application/json'}));
|
| 220 |
+
const a=document.createElement('a');a.href=url;a.download=`${algorithm.toLowerCase()}-${state.job.id.slice(0,8)}.json`;a.click();setTimeout(()=>URL.revokeObjectURL(url),1000);
|
| 221 |
});
|
| 222 |
+
if(PPO){
|
| 223 |
+
$('#g-reset-values').addEventListener('click',()=>{state.values=state.job.rollouts.map(row=>row.tokens.map(()=>null));schedulePreview();});
|
| 224 |
+
$('#g-token-detail').addEventListener('input',event=>{
|
| 225 |
+
if(event.target.id!=='g-value-edit'||!event.target.value.trim()||!event.target.validity.valid)return;
|
| 226 |
+
const [ri,ti]=state.selected;state.values[ri][ti]=Number(event.target.value);schedulePreview();
|
| 227 |
+
});
|
| 228 |
+
$('#g-token-detail').addEventListener('click',event=>{
|
| 229 |
+
if(event.target.id!=='g-reset-value')return;
|
| 230 |
+
const [ri,ti]=state.selected;state.values[ri][ti]=null;schedulePreview();
|
| 231 |
+
});
|
| 232 |
+
}
|
| 233 |
document.addEventListener('keydown',event=>{if((event.metaKey||event.ctrlKey)&&event.key==='Enter'){event.preventDefault();if(!$('#g-sample').disabled)$('#g-form').requestSubmit();}});
|
| 234 |
updateControls();pollStatus();
|
static/index.html
CHANGED
|
@@ -12,7 +12,7 @@
|
|
| 12 |
<body>
|
| 13 |
<header class="site-header">
|
| 14 |
<a class="brand" href="/"><span class="brand-mark">L</span> Layer lab <span class="brand-light">/</span><span class="brand-caption">Layers & attention</span></a>
|
| 15 |
-
<nav><a href="/grpo" class="grpo-nav" title="Open the GRPO sandbox">GRPO</a><a href="#lens-section">Layer explorer</a><a href="#attention-section">Attention</a><a href="#method">Method</a><span id="model-pill" class="model-pill">Qwen3-0.6B</span></nav>
|
| 16 |
</header>
|
| 17 |
<main>
|
| 18 |
<div class="intro">
|
|
|
|
| 12 |
<body>
|
| 13 |
<header class="site-header">
|
| 14 |
<a class="brand" href="/"><span class="brand-mark">L</span> Layer lab <span class="brand-light">/</span><span class="brand-caption">Layers & attention</span></a>
|
| 15 |
+
<nav><a href="/grpo" class="grpo-nav" title="Open the GRPO sandbox">GRPO</a><a href="/ppo" class="grpo-nav" title="Open the PPO sandbox">PPO</a><a href="#lens-section">Layer explorer</a><a href="#attention-section">Attention</a><a href="#method">Method</a><span id="model-pill" class="model-pill">Qwen3-0.6B</span></nav>
|
| 16 |
</header>
|
| 17 |
<main>
|
| 18 |
<div class="intro">
|
static/ppo.html
ADDED
|
@@ -0,0 +1,126 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
<!doctype html>
|
| 2 |
+
<html lang="en">
|
| 3 |
+
<head>
|
| 4 |
+
<meta charset="utf-8">
|
| 5 |
+
<meta name="viewport" content="width=device-width, initial-scale=1">
|
| 6 |
+
<meta name="theme-color" content="#ffffff">
|
| 7 |
+
<title>PPO sandbox — Layer lab</title>
|
| 8 |
+
<link rel="icon" href="data:,">
|
| 9 |
+
<link rel="stylesheet" href="/styles.css">
|
| 10 |
+
<link rel="stylesheet" href="/grpo.css">
|
| 11 |
+
<script type="module" src="/grpo.js"></script>
|
| 12 |
+
</head>
|
| 13 |
+
<body class="grpo-page">
|
| 14 |
+
<header class="site-header">
|
| 15 |
+
<a class="brand" href="/"><span class="brand-mark">L</span> Layer lab <span class="brand-light">/</span><span class="brand-caption">PPO sandbox</span></a>
|
| 16 |
+
<nav><a class="mode-link" href="/">Layers & attention</a><a class="mode-link" href="/grpo">GRPO</a><a class="mode-link active" href="/ppo" aria-current="page">PPO</a><span class="model-pill">Weights frozen</span></nav>
|
| 17 |
+
</header>
|
| 18 |
+
<main>
|
| 19 |
+
<div class="intro">
|
| 20 |
+
<div><p class="eyebrow">PROXIMAL POLICY OPTIMIZATION</p><h1>A value for every state.</h1><p class="intro-description">Sample with a frozen actor. Evaluate each state with a trained value critic.<br>Follow rewards → values → GAE advantages → token gradients.</p></div>
|
| 21 |
+
<div class="g-process" aria-label="Experiment steps"><span><b>01</b> Sample</span><i>→</i><span><b>02</b> Reward</span><i>→</i><span><b>03</b> Gradient</span></div>
|
| 22 |
+
</div>
|
| 23 |
+
<div class="g-workspace">
|
| 24 |
+
<aside class="controls g-controls" aria-label="Rollout settings">
|
| 25 |
+
<div class="sidebar-heading"><h2>Experiment <span class="small-tag">INPUT</span></h2><button id="g-example" class="text-button" type="button">Reset example</button></div>
|
| 26 |
+
<form id="g-form">
|
| 27 |
+
<section class="control-section">
|
| 28 |
+
<label class="field-label" for="g-model">Model</label>
|
| 29 |
+
<select id="g-model" data-sampling><option value="Qwen/Qwen3-0.6B">Qwen3-0.6B</option><option value="ThomasTheMaker/gm3-270m-TinyGSM-reasoning">Gemma 3 270M · TinyGSM reasoning</option></select>
|
| 30 |
+
<p id="g-status" class="field-note" role="status">Connecting to the model…</p>
|
| 31 |
+
</section>
|
| 32 |
+
<section class="control-section">
|
| 33 |
+
<label class="field-label" for="g-mode"><span class="step">01</span> Prefix format</label>
|
| 34 |
+
<select id="g-mode" data-sampling><option value="chat">Chat · official model template</option><option value="raw">Raw · exact tokenized prefix</option></select>
|
| 35 |
+
<label class="field-label g-spaced" for="g-prefix">Prompt / prefix</label>
|
| 36 |
+
<textarea id="g-prefix" rows="4" required maxlength="20000" spellcheck="false" data-sampling>What is 2 + 3? Reply with only the answer.</textarea>
|
| 37 |
+
<div id="g-assistant-field"><label class="field-label g-spaced" for="g-assistant">Assistant prefix <span class="small-tag">OPTIONAL</span></label><textarea id="g-assistant" rows="2" maxlength="5000" spellcheck="false" placeholder="Start the model's answer here…" data-sampling></textarea></div>
|
| 38 |
+
<label class="g-check g-spaced"><input type="checkbox" id="g-thinking" data-sampling> Enable reasoning</label>
|
| 39 |
+
<p id="g-thinking-note" class="field-note">Qwen's native thinking template.</p>
|
| 40 |
+
<details class="g-special"><summary>Insert a control token</summary><div id="g-special-tokens" class="special-tokens"></div></details>
|
| 41 |
+
</section>
|
| 42 |
+
<section class="control-section">
|
| 43 |
+
<label class="field-label" for="g-answer"><span class="step">02</span> Correct answer</label>
|
| 44 |
+
<textarea id="g-answer" rows="2" required maxlength="5000" spellcheck="false">5</textarea>
|
| 45 |
+
<label class="field-label g-spaced" for="g-rule">Automatic reward</label>
|
| 46 |
+
<select id="g-rule"><option value="exact">Exact match · 1 or 0</option><option value="contains">Contains answer · 1 or 0</option></select>
|
| 47 |
+
<p class="field-note">Ignores case and repeated whitespace. Scores the answer after reasoning. Every rollout's reward is editable.</p>
|
| 48 |
+
</section>
|
| 49 |
+
<section class="control-section">
|
| 50 |
+
<div class="g-input-grid">
|
| 51 |
+
<label>Rollouts G<input id="g-count" type="number" value="5" min="2" max="12" step="1" required data-sampling></label>
|
| 52 |
+
<label>Max new tokens<input id="g-length" type="number" value="24" min="1" max="64" step="1" required data-sampling></label>
|
| 53 |
+
<label>Temperature<input id="g-temperature" type="number" value="1" min="0.1" max="2" step="0.1" required data-sampling></label>
|
| 54 |
+
<label>Seed<input id="g-seed" type="number" value="42" min="0" max="4294967295" step="1" required data-sampling></label>
|
| 55 |
+
</div>
|
| 56 |
+
<p class="field-note">Full-vocabulary sampling. Probabilities include temperature; no top-k or top-p filtering.</p>
|
| 57 |
+
<button id="g-sample" class="primary-button" type="submit" disabled><span>Sample 5 rollouts</span><span>↗</span></button>
|
| 58 |
+
<button id="g-cancel" class="g-stop hidden" type="button">Stop sampling</button>
|
| 59 |
+
<p id="g-error" class="g-error hidden" role="alert"></p>
|
| 60 |
+
<p id="g-stale" class="g-warning hidden">Sampling settings changed. Sample a new group to use them.</p>
|
| 61 |
+
</section>
|
| 62 |
+
</form>
|
| 63 |
+
<p class="sidebar-footnote">Real checkpoint samples · no optimizer<br>No model weights are updated.</p>
|
| 64 |
+
</aside>
|
| 65 |
+
<div class="g-results">
|
| 66 |
+
<section class="g-step-panel" aria-labelledby="g-step-title">
|
| 67 |
+
<div class="g-section-heading"><div><p class="eyebrow">03 / THE OBJECTIVE</p><h2 id="g-step-title">Inspect the objective</h2></div><button id="g-export" class="text-button" disabled>Export JSON ↓</button></div>
|
| 68 |
+
<div class="g-ablation-grid g-ppo-ablations">
|
| 69 |
+
<label class="g-ablation"><span><input id="g-normalize" type="checkbox"> Normalize GAE advantages</span><small>Center and scale across all generated tokens</small></label>
|
| 70 |
+
<label class="g-ablation"><span><input id="g-critic" type="checkbox" checked> Use value critic V</span><small>Disable to set all state values to zero</small></label>
|
| 71 |
+
<label class="g-ablation"><span><input id="g-clipping" type="checkbox" checked> Clip the ratio ρᵢ,ₜ</span><small>ρᶜˡⁱᵖᵢ,ₜ = clip(ρᵢ,ₜ, 1−ε, 1+ε)</small></label>
|
| 72 |
+
<label class="g-ablation"><span><input id="g-kl" type="checkbox" checked> KL penalty βKᵢ,ₜ</span><small>Frozen πref = collection πold</small></label>
|
| 73 |
+
</div>
|
| 74 |
+
<p class="g-explanation g-critic-note"><a href="https://huggingface.co/Open-Reasoner-Zero/Open-Reasoner-Zero-Critic-0.5B" target="_blank" rel="noopener noreferrer">Open Reasoner Zero Critic 0.5B</a> · trained scalar value-head. This math critic was trained with a different policy and reward distribution. Its values are illustrative estimates for your prompt, not calibrated success probabilities for the selected actor or custom target. Editing the target does not retrain V. Click a token to override its value.</p>
|
| 75 |
+
<p id="g-critic-status" class="field-note" role="status">Connecting to the value critic…</p>
|
| 76 |
+
<button id="g-reset-values" class="text-button" disabled>Reset all V to checkpoint values</button>
|
| 77 |
+
<div class="g-slider-grid">
|
| 78 |
+
<label>Clip ε <output id="g-epsilon-value">0.20</output><input id="g-epsilon" type="range" min="0.01" max="0.5" step="0.01" value="0.2"></label>
|
| 79 |
+
<label>KL β <output id="g-beta-value">0.04</output><input id="g-beta" type="range" min="0" max="1" step="0.01" value="0.04"></label>
|
| 80 |
+
<label>Discount γ <output id="g-gamma-value">1.00</output><input id="g-gamma" type="range" min="0" max="1" step="0.01" value="1"></label>
|
| 81 |
+
<label>GAE λ <output id="g-lambda-value">0.95</output><input id="g-lambda" type="range" min="0" max="1" step="0.01" value="0.95"></label>
|
| 82 |
+
<label>Candidate drift d <output id="g-drift-value">0.00</output><input id="g-drift" type="range" min="-3" max="3" step="0.05" value="0"></label>
|
| 83 |
+
</div>
|
| 84 |
+
<p id="g-origin-note" class="g-explanation">At d = 0, πθ = πold = πref: ρᵢ,ₜ = 1 and Kᵢ,ₜ = 0. Clipping and KL do not change this first gradient. Move drift in either direction to explore their effects.</p>
|
| 85 |
+
<p id="g-preview-status" class="field-note" role="status">Sample a group to see gradients on every token automatically.</p>
|
| 86 |
+
</section>
|
| 87 |
+
<div class="g-summary g-ppo-summary" aria-live="polite">
|
| 88 |
+
<article><span>Mean reward · mean(r)</span><strong id="g-mean">—</strong><small id="g-std">Terminal reward statistics</small></article>
|
| 89 |
+
<article title="The maximized clipped actor objective, with GAE advantages treated as constants."><span>Objective 𝒥<sub>PPO</sub> ↑</span><strong id="g-loss">—</strong><small>Maximize · advantages from GAE</small></article>
|
| 90 |
+
<article><span>Critic loss L<sub>V</sub> ↓</span><strong id="g-critic-loss">—</strong><small>½(V − detached λ-return)²</small></article>
|
| 91 |
+
<article><span>Clipped tokens</span><strong id="g-clipped-count">—</strong><small id="g-kl-metric">Mean Kᵢ,ₜ: —</small></article>
|
| 92 |
+
</div>
|
| 93 |
+
<section aria-labelledby="g-rollouts-heading">
|
| 94 |
+
<div class="g-section-heading"><div><p class="eyebrow">01 → 02 / SAMPLES & REWARDS</p><h2 id="g-rollouts-heading">The rollout group</h2></div><button id="g-reset-rewards" class="text-button" disabled>Reset auto rewards</button></div>
|
| 95 |
+
<div id="g-progress" class="g-progress hidden" role="status"><span class="spinner"></span><span id="g-progress-text">Sampling…</span></div>
|
| 96 |
+
<p id="g-group-note" class="g-explanation">Choose a prefix and target, then sample a group. No completions are prewritten.</p>
|
| 97 |
+
<p id="g-equal" class="g-warning hidden">Equal rollout rewards do not imply equal PPO advantages; each state has its own value estimate.</p>
|
| 98 |
+
<div class="g-token-legend"><span id="g-token-legend">Each token shows its collection probability and objective gradient.</span><span><i class="g-swatch" aria-hidden="true">+</i> Positive g <i class="g-swatch negative" aria-hidden="true">−</i> Negative g <i class="g-swatch clipped"></i> Clipped</span></div>
|
| 99 |
+
<div id="g-rollouts" class="g-rollouts"><div class="g-empty"><span>G</span><h3>A group makes the comparison.</h3><p>Five real rollouts by default.<br>Every token gets a value, a GAE advantage, and an actor gradient.</p></div></div>
|
| 100 |
+
<details id="g-prefix-details" class="g-prefix-details hidden"><summary>Inspect the exact collection prefix</summary><pre id="g-formatted-prefix"></pre><div id="g-prefix-tokens" class="g-prefix-tokens"></div></details>
|
| 101 |
+
</section>
|
| 102 |
+
<section id="g-inspector" class="g-inspector hidden" aria-labelledby="g-inspector-title"><div class="g-section-heading"><div><p class="eyebrow">SELECTED TOKEN</p><h2 id="g-inspector-title">Follow the gradient</h2></div></div><div id="g-token-detail"></div></section>
|
| 103 |
+
<section class="g-method" id="g-method">
|
| 104 |
+
<p class="eyebrow">WHAT THIS SANDBOX COMPUTES</p><h2>Real samples. Exact objective derivatives.</h2>
|
| 105 |
+
<p>Rollouts and collection probabilities come from the selected actor. A separate frozen ORZ 0.5B critic estimates V(h) before every generated token, using its own tokenizer. Exact histories are retokenized independently; future text is never included. No random value-head, reward-model substitution, or temperature-based pseudo-critic is used.</p>
|
| 106 |
+
<div class="g-equations" aria-label="PPO objective and notation">
|
| 107 |
+
<div class="g-equation-scroll"><math xmlns="http://www.w3.org/1998/Math/MathML" display="block" aria-label="J PPO is the expectation of the group and token averaged minimum of rho times A hat and clipped rho times A hat, minus beta K."><mrow><msub><mi>𝒥</mi><mtext>PPO</mtext></msub><mo>=</mo><mi>𝔼</mi><mo>[</mo><mfrac><mi>1</mi><mi>G</mi></mfrac><munder><mo>∑</mo><mi>i</mi></munder><mfrac><mi>1</mi><mrow><mo>|</mo><msub><mi>y</mi><mi>i</mi></msub><mo>|</mo></mrow></mfrac><munder><mo>∑</mo><mi>t</mi></munder><mrow><mo>(</mo><mrow><mrow><mi mathvariant="normal">min</mi><mrow><mo>(</mo><mrow><msub><mi>ρ</mi><mrow><mi>i</mi><mo>,</mo><mi>t</mi></mrow></msub><msub><mover accent="true"><mi>A</mi><mo>^</mo></mover><mrow><mi>i</mi><mo>,</mo><mi>t</mi></mrow></msub><mo>,</mo><msubsup><mi>ρ</mi><mrow><mi>i</mi><mo>,</mo><mi>t</mi></mrow><mtext>clip</mtext></msubsup><msub><mover accent="true"><mi>A</mi><mo>^</mo></mover><mrow><mi>i</mi><mo>,</mo><mi>t</mi></mrow></msub></mrow><mo>)</mo></mrow></mrow><mo>−</mo><mi>β</mi><msub><mi>K</mi><mrow><mi>i</mi><mo>,</mo><mi>t</mi></mrow></msub></mrow><mo>)</mo></mrow><mo>]</mo></mrow></math></div>
|
| 108 |
+
<div class="g-equation-scroll"><math xmlns="http://www.w3.org/1998/Math/MathML" display="block" aria-label="rho i t is pi theta of y i t given h i t divided by pi old of the same token given the same history."><mrow><msub><mi>ρ</mi><mrow><mi>i</mi><mo>,</mo><mi>t</mi></mrow></msub><mo>=</mo><mfrac><mrow><msub><mi>π</mi><mi>θ</mi></msub><mrow><mo>(</mo><mrow><msub><mi>y</mi><mrow><mi>i</mi><mo>,</mo><mi>t</mi></mrow></msub><mo>|</mo><msub><mi>h</mi><mrow><mi>i</mi><mo>,</mo><mi>t</mi></mrow></msub></mrow><mo>)</mo></mrow></mrow><mrow><msub><mi>π</mi><mtext>old</mtext></msub><mrow><mo>(</mo><mrow><msub><mi>y</mi><mrow><mi>i</mi><mo>,</mo><mi>t</mi></mrow></msub><mo>|</mo><msub><mi>h</mi><mrow><mi>i</mi><mo>,</mo><mi>t</mi></mrow></msub></mrow><mo>)</mo></mrow></mrow></mfrac></mrow></math></div>
|
| 109 |
+
<div class="g-equation-scroll"><math xmlns="http://www.w3.org/1998/Math/MathML" display="block" aria-label="Clipped rho i t is rho i t clipped between one minus epsilon and one plus epsilon."><mrow><msubsup><mi>ρ</mi><mrow><mi>i</mi><mo>,</mo><mi>t</mi></mrow><mtext>clip</mtext></msubsup><mo>=</mo><mrow><mi mathvariant="normal">clip</mi><mrow><mo>(</mo><mrow><msub><mi>ρ</mi><mrow><mi>i</mi><mo>,</mo><mi>t</mi></mrow></msub><mo>,</mo><mi>1</mi><mo>−</mo><mi>ε</mi><mo>,</mo><mi>1</mi><mo>+</mo><mi>ε</mi></mrow><mo>)</mo></mrow></mrow></mrow></math></div>
|
| 110 |
+
<div class="g-equation-scroll"><math xmlns="http://www.w3.org/1998/Math/MathML" display="block" aria-label="TD residual equals immediate reward plus discounted next value minus current value."><mrow><msub><mi>δ</mi><mi>t</mi></msub><mo>=</mo><msub><mi>r</mi><mi>t</mi></msub><mo>+</mo><mi>γ</mi><mi>V</mi><mo>(</mo><msub><mi>h</mi><mrow><mi>t</mi><mo>+</mo><mn>1</mn></mrow></msub><mo>)</mo><mo>−</mo><mi>V</mi><mo>(</mo><msub><mi>h</mi><mi>t</mi></msub><mo>)</mo></mrow></math></div>
|
| 111 |
+
<div class="g-equation-scroll"><math xmlns="http://www.w3.org/1998/Math/MathML" display="block" aria-label="GAE advantage equals TD residual plus gamma lambda times next advantage."><mrow><msub><mover accent="true"><mi>A</mi><mo>^</mo></mover><mi>t</mi></msub><mo>=</mo><msub><mi>δ</mi><mi>t</mi></msub><mo>+</mo><mi>γ</mi><mi>λ</mi><msub><mover accent="true"><mi>A</mi><mo>^</mo></mover><mrow><mi>t</mi><mo>+</mo><mn>1</mn></mrow></msub><mo>,</mo><mspace width="1em"/><msub><mover accent="true"><mi>R</mi><mo>^</mo></mover><mi>t</mi></msub><mo>=</mo><msub><mover accent="true"><mi>A</mi><mo>^</mo></mover><mi>t</mi></msub><mo>+</mo><mi>V</mi><mo>(</mo><msub><mi>h</mi><mi>t</mi></msub><mo>)</mo></mrow></math></div>
|
| 112 |
+
<div class="g-equation-scroll"><math xmlns="http://www.w3.org/1998/Math/MathML" display="block" aria-label="Critic loss is the mean half squared error to a detached lambda-return."><mrow><msub><mi>L</mi><mi>V</mi></msub><mo>=</mo><mtext>mean</mtext><mo>[</mo><mfrac><mn>1</mn><mn>2</mn></mfrac><msup><mrow><mo>(</mo><mi>V</mi><mo>(</mo><msub><mi>h</mi><mi>t</mi></msub><mo>)</mo><mo>−</mo><mtext>stopgrad</mtext><mo>(</mo><msub><mover accent="true"><mi>R</mi><mo>^</mo></mover><mi>t</mi></msub><mo>)</mo><mo>)</mo></mrow><mn>2</mn></msup><mo>]</mo></mrow></math></div>
|
| 113 |
+
</div>
|
| 114 |
+
<p>Both EOS and the maximum token budget end a finite-horizon episode. The final generated action receives the rollout score rᵢ, all earlier rewards are zero, and terminal V = 0. A completion cut off by the token budget is scored as displayed; it is not bootstrapped as an unfinished infinite-horizon trajectory. Stop tokens participate in the objective; prompt tokens do not.</p>
|
| 115 |
+
<p>GAE is computed backward with γ and λ. The actor uses detached advantages; optional normalization centers and scales them across all generated tokens with sample SD + 10⁻⁸. The critic target is the unnormalized λ-return  + V, treated as a constant for its squared-error loss. Both losses use the same group mean of each rollout's token mean. The critic is frozen and can be overridden per state or disabled (V = 0). A scalar value model has no sampling temperature.</p>
|
| 116 |
+
<p>The actor objective uses the clipped ratio ρ = πθ/πold. The optional KL penalty uses frozen πref = πold and K = πref/πθ − log(πref/πθ) − 1. This sandbox adds KL to the actor objective; it does not also add KL reward shaping to GAE.</p>
|
| 117 |
+
<p>Each chip shows the collection probability and <code>gᵢ,ₜ = ∂𝒥PPO/∂log πθ(yᵢ,ₜ | hᵢ,ₜ)</code>, including group/token averaging. Positive g locally favors increasing this log-probability; negative g favors decreasing it. Selecting a token also shows its objective contribution and equivalent loss. The loss is the negative of the maximized objective.</p>
|
| 118 |
+
<p>Candidate drift sets <code>logit(πθ) = logit(πold) + dÂᵢ,ₜ</code>, with d from −3 to +3. This is a synthetic candidate for exploring clipping and KL, not a model forward pass or a measured training change. Positive drift moves with the advantage; negative drift moves against it. At zero drift the candidate is exactly the collection policy.</p>
|
| 119 |
+
<p><strong>A gradient is not a probability change.</strong> Actual training backpropagates <code>∇θ𝒥 = Σ gᵢ,ₜ ∇θ log πθ</code>, updates shared weights with an optimizer, then computes new probabilities with another forward pass. Shared parameters couple tokens, and optimizer state affects the result. This sandbox computes no parameter gradients or optimizer steps and reports no updated-model probabilities. Policy clipping removes an incentive; it does not forcibly clamp probabilities.</p>
|
| 120 |
+
</section>
|
| 121 |
+
<footer class="site-footer"><span>Layer lab <i>/</i> PPO sandbox</span><span>Samples · rewards · objective gradients · frozen weights</span></footer>
|
| 122 |
+
</div>
|
| 123 |
+
</div>
|
| 124 |
+
</main>
|
| 125 |
+
</body>
|
| 126 |
+
</html>
|
static/styles.css
CHANGED
|
@@ -276,6 +276,7 @@ h3{font-size:13px;font-weight:500}
|
|
| 276 |
|
| 277 |
@media(max-width:600px){.site-header{padding:0 16px}
|
| 278 |
.site-header nav a{display:none}
|
|
|
|
| 279 |
main{padding:0 16px}
|
| 280 |
.workspace{display:block}
|
| 281 |
.controls{position:relative;top:0;max-height:none;margin-bottom:24px}
|
|
@@ -315,6 +316,7 @@ main{padding:0 16px}
|
|
| 315 |
}
|
| 316 |
|
| 317 |
@media(max-width:390px){.metrics article:last-child{display:none}
|
|
|
|
| 318 |
.metrics{grid-template-columns:1fr 1fr}
|
| 319 |
.heads-grid{grid-template-columns:repeat(2,minmax(0,1fr))}
|
| 320 |
.matrix-layout{grid-template-columns:1fr}
|
|
|
|
| 276 |
|
| 277 |
@media(max-width:600px){.site-header{padding:0 16px}
|
| 278 |
.site-header nav a{display:none}
|
| 279 |
+
.site-header .brand-light{display:none}
|
| 280 |
main{padding:0 16px}
|
| 281 |
.workspace{display:block}
|
| 282 |
.controls{position:relative;top:0;max-height:none;margin-bottom:24px}
|
|
|
|
| 316 |
}
|
| 317 |
|
| 318 |
@media(max-width:390px){.metrics article:last-child{display:none}
|
| 319 |
+
.site-header .model-pill{display:none}
|
| 320 |
.metrics{grid-template-columns:1fr 1fr}
|
| 321 |
.heads-grid{grid-template-columns:repeat(2,minmax(0,1fr))}
|
| 322 |
.matrix-layout{grid-template-columns:1fr}
|