Flatten GRPO main: Group and Episode

[Jack-Khuu]

The updates boil down to 2 changes that don't alter behavior in grpo/main:

• Group is downgraded from a dataclass to a typedef oflist[Episode], since it's never required
• Episode now directly holds a Completion with redundant attributes in Episode being removed
  • See df0e5a9 for how the redundant fields are mapped between Episode and ScoredCompletion.

(There's also various typehint improvements sprinkled in)

---

Note: This PR does not address or utilize Episode from data_models, but convergence is imminent

python -m apps.grpo.main --config apps/grpo/qwen3_1_7b.yaml

Wandb looks roughly the same
Before: torchforge/grpo-training/runs/wca6wke2
After torchforge/grpo-training/runs/ul34xjr9

[JenniferWang]

• 1 for removing the Group abstraction.

[JenniferWang]

Not related to this diff but now since we're scrutinizing the main flow again, I think making compute_advantages its own Actor is very weird and probably the opposite to an "optimization"

1. We do not expose capability to specify the hostmesh for a specific actor -- ideally, this should be collocated with the generator replica that produces this batch.
2. ComputeAdvantage only needs the rewards; so very likely the entire episodes are serialized.

I wonder, if for now it should be just inlined in the sample call; or allocating a proc on the Policy mesh along side the PolicyWorker to handle the computation but chain the calls and return the result together in sample

[allenwang28]

@JenniferWang these are good points. I want to propose an idea (not for you to implement @Jack-Khuu just brainstorming if this makes sense)

with policy.session() as s:
    host: HostMesh =  await s.get_host_mesh() # returns the host mesh associated with this replica
    advantages = host.run_task(compute_advantages) # where compute_advantages is a function

[joecummings]

Chained calls would be cool 👀

[Jack-Khuu]

This looks legit; +1 on chained calls

[allenwang28]

@JenniferWang these are good points. I want to propose an idea (not for you to implement @Jack-Khuu just brainstorming if this makes sense)

with policy.session() as s:
    host: HostMesh =  await s.get_host_mesh() # returns the host mesh associated with this replica
    advantages = host.run_task(compute_advantages) # where compute_advantages is a function

[joecummings]

Awesome stuff! Just a bunch of small comments.

[joecummings]

Chained calls would be cool 👀