IgorSlinko commited on
Commit
847dfdc
·
verified ·
1 Parent(s): 5fecbd4

Add frozen-critic PPO and show objective gradients directly on tokens

Browse files

Use 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.

Files changed (14) hide show
  1. .dockerignore +2 -0
  2. Dockerfile +2 -2
  3. README.md +24 -6
  4. critic.py +105 -0
  5. grpo.py +65 -38
  6. ppo.py +126 -0
  7. preload_models.py +3 -1
  8. server.py +58 -4
  9. static/grpo.css +11 -3
  10. static/grpo.html +11 -11
  11. static/grpo.js +62 -42
  12. static/index.html +1 -1
  13. static/ppo.html +126 -0
  14. 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 a frozen-model GRPO sandbox. Run it locally or on Hugging Face Spaces. Qwen remains the default model.
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 their tokenizers from Hugging Face; subsequent launches use the standard Hugging Face cache. No API key is required. Qwen loads first, then Gemma, in the background. Each model has its own loading status and remains available after loading. The 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,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
- **Preview one step** differentiates this objective with respect to each sampled-token occurrence's logit, holds all other vocabulary logits fixed, takes an ascent step `z_preview = z_θ + η ∂𝒥_GRPO/∂z`, and exactly renormalizes the categorical distribution. Chips show collection probability → preview probability, with signed changes in **percentage points**. Selecting a token shows the candidate probability, step-only change, ρ_i,t, clipping branch, K_i,t, and ascent-gradient contributions. A large step can overshoot; clipping removes a policy-gradient incentive and does not clamp final probabilities. JSON export includes the actual samples, manual rewards, objective settings, notation, and preview calculations. The API also retains `loss = -objective` and its negative gradients for compatibility; the interface consistently displays the maximized objective.
72
 
73
- **No model parameter gradients, optimizer updates, or trained-model probabilities are computed.** This is an independent-logit educational preview; shared parameters and repeated prefixes do not couple token occurrences. The interface states this distinction explicitly. Both checkpoints stay unchanged. Up to two completed rollout groups per model are retained alongside an active group; older groups and restarts require resampling. Generation shares the model's existing computation lock with the layer explorer.
 
 
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 an explicit, independent-logit GRPO sandbox."""
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
- old = min(max(token["probability"], 1e-12), 1 - 1e-12)
72
- old_logit = math.log(old) - math.log1p(-old)
73
- current = old if payload.drift == 0 else min(max(sigmoid(old_logit + payload.drift * advantage), 1e-12), 1 - 1e-12)
74
- ratio = current / old
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"] != "sampling"]
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 preview."
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"GRPO · {message}"
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"Sampling failed: {exc}"
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"] != "sampling"]
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
- for model_id, spec in MODELS.items():
 
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
- @app.post("/api/grpo/runs")
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 grpo_jobs.start(runtime, payload, context, ids)
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-slider-grid{display:grid;grid-template-columns:repeat(4,minmax(0,1fr));gap:20px;margin-top:19px}
 
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-delta{font:600 11px var(--mono);white-space:nowrap}
 
 
 
 
 
 
77
  .g-token.up{background:#fff;border-color:#aaa}.g-token.down{background:#fff;border-color:#aaa}
78
- .g-token.up .g-token-delta{color:#222}.g-token.down .g-token-delta{color:#222}
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. One step.</h1><p class="intro-description">Sample real completions, assign rewards, and follow the gradient.<br>Explore a step in logit space while the model stays fixed.</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> Preview</span></div>
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">Build the step</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,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
- <div class="g-preview-action"><button id="g-preview" class="primary-button" disabled><span>Preview one step</span><span>→</span></button><button id="g-hide-preview" class="text-button hidden">Hide probability changes</button><span id="g-preview-status" class="field-note" role="status">Sample a group first.</span></div>
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> Increase <i class="g-swatch negative" aria-hidden="true">−</i> Decrease <i class="g-swatch clipped"></i> Clipped</span></div>
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 PREVIEW COMPUTES</p><h2>A real objective. An illustrative probability step.</h2>
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>Candidate drift sets <code>logit(πθ) = logit(πold) + dÂᵢ</code>, with d from −3 to +3. Positive drift moves with the advantage; negative drift moves against it. The preview takes one exact gradient-ascent step on 𝒥GRPO for each occurrence's sampled-token logit, holds all other vocabulary logits fixed, then renormalizes. Step size η belongs to this illustration and is not a model learning rate.</p>
108
- <p><strong>No model parameter gradients or optimizer steps are computed.</strong> Token occurrences are independent in this sandbox; it does not simulate how shared network parameters couple tokens or repeated prefixes. The displayed probabilities are a local logit-space preview, not predictions from a trained model. Neither ratios nor final probabilities are forcibly clamped by GRPO clipping.</p>
 
109
  </section>
110
- <footer class="site-footer"><span>Layer lab <i>/</i> GRPO sandbox</span><span>One group · one hypothetical step · frozen weights</span></footer>
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 state = {job:null, preview:null, rewards:[], showStep:false, selected:null, status:null, busy:false, previewBusy:false};
 
 
 
 
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:Number($('#g-step').value), drift:Number($('#g-drift').value)};
 
27
  }
28
  function updateControls(){
29
- const sampling=state.busy||state.job?.phase==='sampling', ready=state.job?.phase==='ready';
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?.phase!=='sampling');
39
- $('#g-preview').disabled=sampling||!ready||state.previewBusy||!state.preview;
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=state.showStep?'Collected πold → preview πpreview · Δ in percentage points vs collected':'Each token shows πold(yᵢ,ₜ | hᵢ,ₜ), its collection probability.';
48
- for(const [id,digits] of [['epsilon',2],['beta',2],['step',1]])$('#g-'+id+'-value').textContent=Number($('#g-'+id).value).toFixed(digits);
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
- 'Candidate drift shifts each sampled-token logit by d × Âᵢ: positive d moves with the advantage; negative d moves against it. Displayed Δ includes both drift and gradient ascent on 𝒥GRPO. Click a token to separate them.';
 
 
 
 
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==='sampling'||state.status?.phase!=='ready'?1800:5000);
74
  }
75
 
76
  function tokenMarkup(token, ri, ti){
77
- const detail=state.preview?.rollouts[ri]?.tokens[ti], show=state.showStep&&detail;
78
- const change=show?detail.delta:0;
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?`\nCurrent πθ: ${pct(detail.current)}\nPreview πpreview: ${pct(detail.after)}\nTotal change: ${signed(detail.delta*100,6)} pp\nStep alone: ${signed(detail.step_delta*100,6)} pp\nRatio ρᵢ,ₜ: ${number(detail.ratio)} · ${detail.clipped?'policy gradient clipped':'policy gradient active'}`:''}`;
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-after="${detail.after}" data-delta="${detail.delta}" data-clipped="${detail.clipped}"`:''} title="${escape(title)}" aria-label="${escape(title)}">${show?`<span class="g-token-delta">${signed(change*100)} pp</span>`:''}<span class="g-token-token">${escape(visible(token.text))}</span><span class="g-token-prob">${pct(token.probability)}${show?` → ${pct(detail.after)}`:''}</span></button>`;
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.phase==='sampling';
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=state.showStep?'Logit-space preview · weights unchanged':'Rewards ready. Preview the hypothetical step.';
 
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)],['Candidate πθ',pct(detail.current)],['Preview πpreview',pct(detail.after)],['Step alone',signed(detail.step_delta*100,4)+' pp'],['Advantage Âᵢ',signed(row.advantage,4)],['Policy ratio ρᵢ,ₜ',number(detail.ratio)],['Sampled KL Kᵢ,ₜ',number(detail.kl)],['Policy branch',detail.clipped?'Clipped · zero slope':'Unclipped']];
124
- $('#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">Token objective, before averaging: policy contribution to the log-probability gradient <code>${number(detail.objective_policy_gradient)}</code> · KL contribution <code>${number(detail.objective_kl_gradient)}</code>.<br>After the softmax derivative and group/token averaging: <code>∂𝒥GRPO/∂z = ${number(detail.objective_logit_gradient)}</code> (averaging weight <code>${detail.weight.toPrecision(4)}</code>).<br>Gradient ascent: <code>z_preview = z_θ + η × ∂𝒥GRPO/∂z</code>. Other vocabulary logits stay fixed; softmax renormalizes the entire distribution.${detail.clipped?' The policy term has zero slope here. An enabled KL term can still move the token.':''}</p>`;
 
 
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(`/api/grpo/runs/${sid}/preview`,objectivePayload());
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='Preview unavailable. Check the reward settings.';}}
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(`/api/grpo/runs/${sid}`);
146
  if(state.job?.id!==sid)return;
147
  state.job=job;renderJob();
148
- if(job.phase==='sampling')jobTimer=setTimeout(()=>pollJob(sid),800);
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?.phase==='sampling')return;
158
  if(!$('#g-form').reportValidity())return;
159
  const payload=collectionPayload();
160
  error();state.busy=true;updateControls();
161
  try{
162
- const job=await api('/api/grpo/runs',payload);
163
  clearTimeout(jobTimer);clearTimeout(previewTimer);previewVersion++;
164
- state.job=job;state.preview=null;state.rewards=[];state.selected=null;state.showStep=false;state.previewBusy=false;
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(`/api/grpo/runs/${state.job.id}/cancel`,{});}catch(exc){error(exc.message);}
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-step','g-drift'])$('#'+id).addEventListener('input',()=>{updateControls();schedulePreview();});
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?.phase==='sampling')return;
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:'grpo-logit-space-preview',weights_updated:false,collection:state.job,objective:objectivePayload(),preview:state.preview,preview_visible:state.showStep,
208
- notation:{objective:'J_GRPO (maximized)',reward:'r_i',advantage:'A_hat_i',ratio:'rho_i_t',clipped_ratio:'rho_i_t^clip',kl:'K_i_t',current:'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_GRPO (minimized)'},
209
- note:'Real frozen-policy samples. Independent sampled-token logit ascent on J_GRPO; no model parameter gradients, optimizer updates, or parameter-coupled predictions.'};
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=`grpo-${state.job.id.slice(0,8)}.json`;a.click();setTimeout(()=>URL.revokeObjectURL(url),1000);
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}