AgentWorkflowPPOTrainer: GRPO groups by trajectory.uid instead of prompt/agent key

[TtLuckyyy]

Question

I'm not sure whether the current behavior is intentional or I'm misunderstanding the design. Reporting what I observed vs. what I expected — would appreciate maintainer confirmation.

Context

Expected behavior (as I understand it)

Config	GRPO grouping uid	Are 4 rollouts grouped as one trajectory-level group?
stepwise_advantage.enable=False	Prompt-level UUID (task_ids)	Yes
mode=per_step	{prompt_uuid}_step{k}	No (group by step index within the same prompt)
mode=broadcast	Prompt-level UUID; compute advantage on last step only, then broadcast to earlier steps	Yes (trajectory-level comparison, then copy to each step)

Actual behavior (current code)

Before compute_advantage, all three configs do:

batch.non_tensor_batch["uid"] = batch.non_tensor_batch["step_ids"]

step_ids comes from trajectory.uid in rllm/experimental/verl/transform.py (a unique UUID per trajectory), which does not match the table above:

Config	Current GRPO uid	Are 4 rollouts grouped together?
enable=False	trajectory.uid (4 distinct values)	No (group size 1 → advantage ≈ raw reward)
mode=per_step	Still trajectory.uid	No (neither prompt-level nor {uuid}_step{k})
mode=broadcast	Still trajectory.uid	No; also, transform emits all rows with is_last=True, so the broadcast path appears largely ineffective

stepwise_advantage currently only changes the reward source (step_rewards vs traj_rewards) and the broadcast flow — it does not change the GRPO grouping key.

Possible impact

If prompt-level GRPO is intended, multiple rollouts of the same prompt may not be compared within a group. The unified trainer path uses TrajectoryGroup for advantage computation and may not be affected.

Suggested fix (if this is a bug)

• enable=False / mode=broadcast: use task_ids (single agent) or trajectory_ids ({task_id}_{name}, multi-agent) as GRPO uid
• mode=per_step: use {task_id}_step{k} (or whatever step index format transform uses)
• broadcast: may also require transform to emit rows with is_last_step=False for intermediate steps

Minimal repro

1 prompt, rollout.n=4, rewards [1, 0, 1, 0]:

• Current: advantages ≈ [1, 0, 1, 0]
• Expected (if enable=False): group mean = 0.5 → advantages ≈ [+0.87, -0.87, +0.87, -0.87] with norm_adv_by_std_in_grpo=true

Files: rllm/trainer/verl/agent_workflow_trainer.py, agent_workflow_trainer_fireworks.py, rllm/experimental/verl/transform.py

Relevant Code / Config

Environment

No response