mzio/aprm-thought-generations-tau2-airline

Act-PRM Thought Generations — tau2-bench Airline

Latent thoughts inferred behind logged, action-only agent demonstrations by the Act-PRM offline EM (Action Process Reward Models), for the tau2-bench airline domain. For each logged (state s, action x) the EM-trained generator samples G=4 candidate thoughts z, and each candidate is scored by the length-penalized action likelihood of the logged action:

likelihood = p(x | s, z)                       # policy per-action-token likelihood
reward     = likelihood - 0.15 * (thought_len / 96)   # reward_method="penalty"
best       = argmax(reward)                     # top-1 selection

Generator = Qwen3-4B-Instruct-2507 LoRA. Generation params: group_size=4, length_penalty=0.15, max_thought_tokens=96, thought_temperature=1.0, reward_method="penalty".

Contents

generations/ — the full candidate pool (one JSONL row per logged action-step). This is the primary artifact: it contains all G thoughts with their scores, so you can select top-1 or recompute any weighting (em / top_half / grpo / uniform) yourself.

Row schema:

{
  "split": "train"|"eval", "sample_id": int, "timestep": int, "batch_id": int, "try_step": int,
  "advantage_mode": "best",
  "target_action": str,          # the logged (ground-truth) action being explained
  "thoughts":      [str × G],    # the G candidate thoughts
  "likelihoods":   [float × G],  # p(x | s, z) per candidate
  "rewards":       [float × G],  # likelihood - 0.15*len_frac per candidate
  "advantages":    [float × G],  # per advantage_mode (here "best" -> one-hot on best)
  "thought_tokens":[int × G],    # candidate thought length in tokens
  "best": int                    # argmax(rewards) — the top-1 index
}

Selecting thoughts

• —Top-1 (what the paper's SFT uses): thoughts[best].
• —EM posterior: softmax-style group-normalized weights over rewards/likelihoods.
• —top_half / grpo / uniform: derive from rewards (grpo: reward - mean, opt. ÷ std).

sft_corpus_top1/ — convenience: the baked top-1 Stage-2 SFT corpus in chat-messages format (assistant content = thoughts[best] + logged action), split into train.json / eval.json per variant. Equivalent to taking thoughts[best] from generations/ and joining onto the logged trajectories.

Splits

21 train trajectories / 4 eval trajectories (the rl_eval hold-out is carved out separately for the RL phase). 258 action-step rows per variant.

Related

• —SFT LoRA adapters: `mzio/aprm-sft-tau2-airline`
• —Source expert demos: `mzio/aprm-tau2-airline-gpt5m_med-gs4-s0-train`