Skip to content
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
5 changes: 5 additions & 0 deletions verifiers/v1/graph.py
Original file line number Diff line number Diff line change
Expand Up @@ -101,6 +101,11 @@ class MessageNode(BaseModel):
logprobs: list[float] = Field(default_factory=list)
"""Sampling logprobs for the sampled tokens — length equals the number of True entries in
`mask`; empty for input messages."""
advantages: list[float] | None = None
"""Per-token credit over the sampled tokens, same layout as `logprobs`. `None` until a
consumer's RL algorithm assigns it, which is not the same as a credit of zero: a group whose
rewards were all equal is assigned zeros and carries no gradient, while an unassigned node was
never scored at all."""
multi_modal_data: SkipJsonSchema[MultiModalData | None] = None
"""The renderer items for the images this message's content introduces (pixel tensors,
grids, hashes, placeholders) — the only carrier of the pixels from the env server to the
Expand Down
41 changes: 29 additions & 12 deletions verifiers/v1/trace.py
Original file line number Diff line number Diff line change
Expand Up @@ -3,7 +3,7 @@
import time
import traceback
import uuid
from collections.abc import Mapping
from collections.abc import Callable, Mapping
from typing import TYPE_CHECKING, Annotated, Any, Generic, Literal

import numpy as np
Expand Down Expand Up @@ -200,28 +200,45 @@ def sampled_mask(self) -> list[bool]:
mask.extend(node.mask)
return mask

@property
def logprobs(self) -> list[float]:
"""Per-token sampling logprobs aligned to `token_ids` — the node logprobs spread onto
their sampled positions, 0.0 on every non-sampled token."""
def spread(
self, values: Callable[[MessageNode], list[float] | None]
) -> list[float]:
"""A per-sampled-token node field widened to `token_ids`: each node's values land on its
sampled positions, 0.0 everywhere else. A node holding nothing contributes zeros."""
out: list[float] = []
for node in self.nodes:
mask = node.mask
sampled = sum(mask) if node.logprobs else 0
mask, node_values = node.mask, values(node) or []
sampled = sum(mask) if node_values else 0
# Bulk-fill the canonical unsampled-prefix/sampled-suffix layout.
if not sampled or all(mask[-sampled:]):
out += [0.0] * (len(mask) - sampled) + node.logprobs[:sampled]
out += [0.0] * max(0, sampled - len(node.logprobs))
out += [0.0] * (len(mask) - sampled) + node_values[:sampled]
out += [0.0] * max(0, sampled - len(node_values))
continue
li = 0
for sampled in mask:
if sampled:
out.append(node.logprobs[li] if li < len(node.logprobs) else 0.0)
for is_sampled in mask:
if is_sampled:
out.append(node_values[li] if li < len(node_values) else 0.0)
li += 1
else:
out.append(0.0)
return out

@property
def logprobs(self) -> list[float]:
"""Per-token sampling logprobs aligned to `token_ids` — the node logprobs spread onto
their sampled positions, 0.0 on every non-sampled token."""
return self.spread(lambda node: node.logprobs)

@property
def advantages(self) -> list[float] | None:
"""Per-token credit aligned to `token_ids`, spread like `logprobs` — or `None` when no node
on the path was ever assigned any, which a branch of zeros would otherwise be
indistinguishable from. A partially assigned path spreads, the unassigned nodes reading
0.0."""
if all(node.advantages is None for node in self.nodes):
return None
return self.spread(lambda node: node.advantages)

@property
def multi_modal_data(self) -> MultiModalData | None:
"""Node image data concatenated in token order for training; never persisted."""
Expand Down
Loading