mzio/aprm-sft-thoughts-tau2-retail

Act-PRM SFT thoughts — tau2-bench retail

Act-PRM (Action Process Reward Models) infers the latent thoughts behind logged, action-only agent demonstrations via an offline EM. For each logged action x in state s we sample G=4 candidate thoughts z, score each by the length-penalized action likelihood

reward(z) = p(x | s, z) - 0.15 * len_frac

(len_frac grows with the thought's token length), and mark the best thought (argmax reward). The (thought + action) span is then what downstream SFT / RL trains on.

This dataset publishes the full E-step: every one of the G sampled thoughts per logged action, each with its reward, likelihood (p(x|s,z)) and thought_tokens, plus the best_index. That is, it is not reduced to the top-1 — you can re-derive top-1, top-half, EM-weighted, or any other selection downstream (see below).

Variants (files / splits)

Four relabel passes = 2 scorers x 2 EM checkpoints:

• —scorer — which model computed p(x|s,z) in the E-step: the trained policy LoRA, or the frozen base model (--score_with_base).
• —em_checkpoint — which EM training checkpoint the scorer was resumed from: the best-metric step, or the last step.

Every row also carries scorer / em_checkpoint columns, so you can also concatenate all four files and filter.

Note: train-split rows have the full G=4 sampled thoughts; the held-out eval-split rows were relabeled with a single sample (G=1), so their thoughts/rewards lists have length 1 and best_index == 0.

Schema (one row per relabel-variant x logged action-step)

Selecting thoughts

from datasets import load_dataset
ds = load_dataset("mzio/aprm-sft-thoughts-tau2-retail", split="policy_best")   # or base_best / policy_last / base_last
row = ds[0]

# top-1 (what the paper's SFT commits): the best length-penalized thought
top1 = row["thoughts"][row["best_index"]]
assert row["rewards"][row["best_index"]] == max(row["rewards"])

# the SFT target span is thought + action:
sft_target = top1 + "\n\n" + row["target_action"]

# top-half: keep the thoughts whose reward is in the top 50%
import numpy as np
order = np.argsort(row["rewards"])[::-1]
top_half = [row["thoughts"][i] for i in order[: max(1, len(order) // 2)]]

# EM weights: group-normalized (softmax-like) weights over all G thoughts,
# e.g. a temperature-1 softmax over rewards (or normalize exp(likelihood)):
import math
r = row["rewards"]
Z = sum(math.exp(x) for x in r)
em_weights = [math.exp(x) / Z for x in r]   # weight every (thought+action) span

Provenance

Generated by Act-PRM Stage-1 relabel passes over tau2-bench retail expert demonstrations (no_train generate-only passes), model hf_qwen3_4b_instruct, group_size=4, length_penalty=0.15. The held-out RL-eval tasks are excluded. Built with scripts/export_sft_dataset_hf.py.