Support getting each turn in a multi-turn task as single training sample

[LuanBrt]

I’d like to request a feature to support splitting a n-turn trajectory into n training samples as described in the Kevin paper, as a way to improve sample efficiency.

As described in the paper:
In each multi-turn training step:

1. For each task, we sample m parallel trajectories with n refinement turns. To improve sample efficiency, each refinement turn (CoT + response) in a trajectory becomes a single training sample. The response of the model after the CoT consists of a kernel and a CoT summary.
2. We construct the context of a sample by including the history of previous responses, which include generated kernels along with their summarized CoTs, and evaluation feedback.
3. We evaluate the generated kernel and compute its score as shown in Section 3.2. The reward of each turn (CoT + response) is the discounted sum of current and subsequent scores, which we elaborate in Section 4.3.
4. For each task, we normalize the rewards across the mn samples for advantage calculation. Then we compute the GRPO loss over the entire batch.

This feature could also be used in order to summarize previous steps of the trajectory allowing for better management of context size, where the LLM would be able to see just the env response of the previous step.

[willccbb]

Agree that this is desirable eventually, but it basically requires you to use a different algorithm than GRPO (this would essentially become REINFORCE), and there's not a lot of writing out there about the implementation details needed to get it to work well for multi-turn agent training (the Kevin paper is a bit vague on things like hyperparameters). It's a research project, and likely a standalone algorithm implementation rather than just a GRPO tweak, happy to review a PR if anyone wants to give it a swing + do some experiments.