[RL] - MessageEnv, Rollout types, Rubric, Renderer

[felipemello1]

How to review?
a. Read the contents in tasks/sum_digits
b. Read grpo.py:collect_rollouts
c. Read the rest

Summary

1. A message-level env protocol:

Our current script does not use messages or chat template. Now it will be the default. Users write reset / step_message;

class SumDigitsEnv(MessageEnv):
    async def reset(self) -> MsgResponseReset:
		...
    async def step_message(self, msg: Message) -> MsgResponseStep:
        ...

a RendererEnv wraps it and owns all message <-> token plumbing done by the Renderer.

example = dataset.get_sample()
env = RendererEnv( 
	message_env=SumDigitsEnv(env_input=example.env_input),
    renderer=renderer
)
initial_turn = env.reset()

2. Typed rollout records: RolloutGroup(Rollout(RolloutTurn)) replace the old Trajectory / (Completion, Step) pairs. They now carry messages and tokens that support multi-turn.

3. Rubric: Class to hold functions for scoring after rollout is finished

class MyRubric(Rubric):
    def register_funcs(self) -> list[RewardFn]:
        return [
            RewardFn(fn=my_reward_fn1, weight=0.5),
            RewardFn(fn=my_reward_fn2, weight=0.5),
        ]

rubric = MyRubric(config=MyRubric.Config(truncation_reward=0.0))
rewards = await rubric.score_group(my_rollouts, env_input)
for reward, rollout in zip(rewards, my_rollouts):
    my_rollout.reward = reward.reward
    my_rollout.reward_components = reward.components

It also handles partial scoring in case of truncation and error.

4. Task: A class that knows how to

a) create/store Envs for an specific task
b) Holds the rubric associated to that task
c) in the future will hold the rollout loop -- which can be customized by users if they want to.

This makes it trivial to do i) dataset mix; ii) share/import tasks

experiments/rl/
├── grpo.py                  # controller + rollout loop  (changed)
├── renderer.py              # NEW: RendererConfig -> a renderers.Renderer
├── env_types/                    # NEW: the env protocol
│   ├── message_env.py       #   MessageEnv (ABC), ResetOutput, StepOutput
│   └── renderer_env.py      #   RendererEnv, RendererEnvConfig, TokenizedStepOutput
├── rollouts/                # NEW: the datatypes
│   ├── types.py             #   Rollout, RolloutTurn, RolloutGroup, RolloutStatus, DatasetOutput
│   └── utils.py             #   rollout_to_episode, prepare_rollout_metrics
├── rubrics/rubric.py        # NEW: Rubric, RewardFn, Reward
└── tasks/                   # NEW (was sum_digits.py)
    ├── task.py              #   Task
    └── sum_digits/          #   the worked example
        └── data.py · env.py · grader.py · task.py

Why?

For single turn we can naively concatenate prompt+response. To enable multiturn we had to enable (1) and (2). Given i was refactoring it, i added (3) and (4)

Blockers/next steps:

1. Continuous Batching (CB): Necessary so multiturn Rollouts can progress independently
2. Async rollout + AlphabetSort multiturn task: Small PR once this and continuous batching lands

[tianyu-l]

Thanks, did a first pass. Overall I feel the logic among Env, Task, Rollout, Dataset, Rubric could be more clear.

[tianyu-l]

nit: TrainingSample, to be consistent with TrainingBatch

[tianyu-l]

would this be used for sft as well?

[wwwjn]

Do we need to depend on renders lastest main?

[tianyu-l]

Can we avoid this dependency?

[felipemello1]

soon: PrimeIntellect-ai/renderers#70

[tianyu-l]

Not clear what this means. We should make it explicitly referring to torchtitan model (and their tokenizer) and avoid transformers dependency.

[tianyu-l]

Suggested change

	stop_token_ids=list(_sampling_config.stop_token_ids) or None,
	stop_token_ids=_sampling_config.stop_token_ids or None,

[tianyu-l]

can we call it reward which should always decompose into "components"

[felipemello1]

I think that we want to have a simple field "reward: float", because thats what the final loss will use. No ambiguity here. And components is the breakdown for logging or weight averaging or some specific advantage computation.

[tianyu-l]

If single field reward is always coming from the components, can we have cached_property that derive the single field from components. If not (instead single field is the only thing RL algorithm eventually care, and everything else is only for logging), then I'm fine with it.

[tianyu-l]

how is env_messages related to prompt_messages?

[felipemello1]

Prompt_messages: input to the generator (all history up to that point)
assistant_message: output of genereator
env_message: output of the environment, e.g. tool calls, new user message, etc

[tianyu-l]

butsomewhere else you used next_prompt_...

Is it the same as env_message or a subset?

[tianyu-l]

so "prompt" is previous [env + assistant] history, but next_prompt only means partial data in a turn

assistant vs. generator used interchangeably?

[tianyu-l]

"Output" sounds confusing, it can be input to rollout / grader

[felipemello1]

maybe DataSample?

[tianyu-l]

The controller can just do: workflow.run_rollout, and it will get the rollout, without knowing about it's internals.

I feel we should do this. The RLTrainer should just own trainer, generator, define task / env and maybe pass generator to them to obtain rollouts.

[wwwjn]

After this PR we are still in Sync RL, but defining these functions here as they are basic for async RL?

[felipemello1]

we do asyncio.gather because the reward_fns can be multiple LLMs, for example. So we can run them in parallel. This is orthogonal to async/sync RL. Does this make sense?

[wwwjn]

What's this helper function for? Can we just inline it?

[wwwjn]

We should set up CPU CI test for RL to guard these tests

[wwwjn]

What's the consideration of putting the dataset a field of SumDigitsTask.Config() vs. connecting them using TASK_NAME? A Task can be the container of Env, Rubric and Dataset?

[wwwjn]

Can we rename to rubrics.py which more aligned with our naming now

[wwwjn]

We only have one task in our current RL loop now, are you going to support data mix soon? I guess we can simplify in this PR and only support one task for now, and handle multi-tasks together with data-mix PR

[wwwjn]

Is stop_token_ids same as eos_ids?

[felipemello1]

thats my understanding, but we should get it directly from the renderer.tokenizer. Also, maybe the user can have some specific logic to stop when some token T appears

[wwwjn]

So the reason we don't put dataset a subfield of Task is because of data-mixing? We will do datamixing at DatasetOutput level, not Task level? What does other repo model these concepts?

[felipemello1]

What does other repo model these concepts?

I dont recall. I can take a look. Another option is to play it by year, so what we are comfortable with and refactor later when we try datamix. For now, lets put it inside of the Task, since thats a pattern i have seen as well and both and tianyu shared that it made more sense to you. I will make the changes.

[wwwjn]

Do we need to create enviornments repeatedly for each sample? Eg, spin up a docker for each single CodingTask?

Or can we create a Enviornment for each Task?

[felipemello1]

its up to the user to decide what happens inside of "make_envs", i.e. create a fresh new one or pull from a pool

[wwwjn]

What is initial_prompt here? Can you give an example? I'm confused why the prompt doesn't come from dataset

[felipemello1]

the sample (history of messages) comes from the dataset. The env adds the system message and adds tool calls. The prompt is the final constructed input for the generator. But notice that we are not opinionated about it: if the user wants the env_input to include the system prompt as well, they can do it.

[tianyu-l]

Overall structure looks reasonable. I still ranted a lot on variable naming, as I think that's the key to hackability.

[tianyu-l]

In general I would prefer full spelling over shorthand.

Could we do

• rl/environments/message.py
• rl/environments/renderer.py

[felipemello1]

I wanted to signal somehow that those are "base" envs, not something like a SumDigitsEnvs, thats why i put env_types. I am cool with environments, but i still think that we should signal somehow that those are types. Any ideas?

[tianyu-l]

I see. If we follow strictly the sft counterpart, for classes / interfaces you would expect users to implement, it probably should be

• protocols/environments
• protocols/rubrics
• etc.

Meanwhile, components for things that torchtitan implement for you (but also overridable).

[tianyu-l]

Can we call this MessageInitialOutput or MessageInitOutput.
reset sounds cleaning up from an old state which this message shouldn't care.

If you want strong consistency with env.reset, I proposal change that as well to env.init

[tianyu-l]

It also sounds unnatural when MessageOutput owns messages and tools.

Maybe should call RawOutput / EnvOutput / UnrenderedOutput vs. RenderedOutput

[tianyu-l]

Reserving prompt for the first message sounds ambiguous.

Can we call it init_messages or just messages?

[tianyu-l]

this check should be added for init output as well?

[tianyu-l]

is this too controversal

Suggested change

	async def reset(self) -> MessageResetOutput:
	async def init(self) -> MessageResetOutput:

[tianyu-l]

"steps" not clear

[tianyu-l]

shouldn't use get()?

[tianyu-l]

What's the consideration for these two phases. I understand that each rollout may finish in different number of turns -- how would we deal with the inefficiency of "finish one rollout at a time" later?

[felipemello1]

when we land the generator refactor, there wont be a "batch" concept. We will change this function to continuously feed new examples to a queue that produces RolloutGroup in its own stream and put in the buffer. does this answer the question?

[tianyu-l]

Suggested change

	num_prompts_per_step: int = 5
	num_groups_per_iteration: int = 5

[tianyu-l]

Add comment on how this number is derived. It looks you are using asyncio at multiple places, does this number capture all the concurrency?

[felipemello1]

@tianyu-l , i saw that you added more comments. I will read and address them.

---

Regarding my latest changes:

Thanks for the thorough review @tianyu-l, @wwwjn. I pushed some changes. Here is the before/after:

Goal for the code (we cannot do it until we fix the generator)

async def do_single_rollout(sampler, env) -> Rollout:
    turns = []
    prev_step = await env.reset()                           # TokenizedStepOutput (first prompt)
    while not prev_step.status.is_terminal():
        completion = await sampler.generate(prev_step.next_prompt_token_ids)
        step = await env.step(completion)
        turns.append(
            RolloutTurn(
                prompt_token_ids=prev_step.next_prompt_token_ids,    # input  -> prompt_*
                prompt_messages=prev_step.next_prompt_messages,
                assistant_token_ids=completion.token_ids,           # model  -> assistant_*
                assistant_logprobs=completion.token_logprobs,
                assistant_message=step.assistant_message,
                env_messages=step.env_messages,                     # env    -> env_*
                env_rewards=step.env_rewards,
            )
        )
        prev_step = step
    return Rollout(turns=turns, status=prev_step.status)

async def do_group_rollout(task, task, group_size) -> RolloutGroup:
    example = task.sample_train_example()
    envs = task.make_envs(group_size=group_size)
    for env in envs:
        rollout = do_single_rollout(env, example)
    rollout_group = RolloutGroup(rollouts=rollouts)
    rewards = task.score_group(rollout_group)
    ...

The overall picture — what a Task bundles, and who talks to whom:

Task — one config bundle
   ├─ train / validation datasets
   ├─ make_envs(env_input)
   ├─ rubric (RewardFn, …)
   └─ #TODO hold the rollout loop after generation fix

A rollout — the user side speaks messages, the model side speaks tokens:

   MessageEnv   ◄── messages ──►   RendererWrapperEnv   ◄── tokens ──►   generator (vLLM, n=1)

      • MessageEnv          user code, speaks messages, never sees tokens
      • RendererWrapperEnv  converts messages ⇄ tokens, and owns the
                            prompt-too-long / truncation / timeout checks

---

1. The env is two explicit layers: messages vs tokens

"step_completion sounds confusing" · "MessageEnv.step_message(msg) sounds confusing — is msg always from assistant?" · "what does 'a Renderer for tokens' mean" · "components sounds confusing, does rewards work?"

Change: The old RendererEnv read like an RL env, but it is a token-space wrapper around the user's message-space env. We named the two layers and gave both the same reset() / step() shape. Every model-output field is unified on assistant_*, every env field on env_*.

Before

class MessageEnv(abc.ABC):
    async def reset(self) -> ResetOutput: ...
    async def step_message(self, msg: Message) -> StepOutput: ...

class RendererEnv:                                    # "an Env that renders"?
    async def initial_prompt(self) -> TokenizedStepOutput: ...
    async def step_completion(self, completion) -> TokenizedStepOutput: ...

class RolloutTurn:
    prompt_token_ids: list[int]
    response_token_ids: list[int]
    response_logprobs: list[float]
    policy_version: int
    prompt_messages: list[Message]
    assistant_message: Message | None
    env_messages: list[Message]
    reward_components: dict[str, float]

After

class MessageEnv(Configurable, abc.ABC):             # user code: messages only, never tokens
    async def reset(self) -> MessageResetOutput: ...
    async def step(self, assistant_message: Message) -> MessageStepOutput: ...

class RendererWrapperEnv:                             # wraps a MessageEnv: messages <-> tokens via a Renderer
    async def reset(self) -> TokenizedStepOutput: ...
    async def step(self, completion: Completion) -> TokenizedStepOutput: ...

# full snapshot of tokens + messages
class RolloutTurn:                               # the full per-turn snapshot
    prompt_token_ids: list[int]                  # input
    prompt_messages: list[Message]
    assistant_token_ids: list[int]               # model output  (was response_token_ids)
    assistant_logprobs: list[float]              # (was response_logprobs)
    assistant_message: Message | None
    env_messages: list[Message]                  # env reply
    env_rewards: dict[str, float]                # (was reward_components)
    policy_version: int | None

The wrapper does tokens → messages (parse the completion) and messages → tokens (render the next prompt), and owns the "prompt too long / generation truncated / env timed out" checks, so the user's MessageEnv only ever sees a clean assistant message:

2. The Task owns its dataset, env, and rubric (config-driven)

"why this env doesn't own the sourcing of prompts / data?" · "why this is not RendererEnv.Config and owned by that class?" · "The relationship between Task and Env seems not clear" · "RLTrainer should just own trainer, generator, define task / env"

Change: A Task is now one config bundle — datasets, env, rubric, env-wrapper limits. The base builds everything and make_envs is generic. A concrete task is just typed defaults.

Before — datasets on the trainer, Task base owned only the rubric, a task_name → task map routed rows:

# RLTrainer.Config
train_dataset = SumDigitsDataset.Config(seed=42)
validation_dataset = SumDigitsDataset.Config(seed=99)
tasks = {"sum_digits": SumDigitsTask.Config()}

class Task(Configurable, abc.ABC):
    class Config:
        rubric: Rubric.Config                        # base owns ONLY the rubric
    @abc.abstractmethod
    def make_envs(self, *, example, group_size, renderer): ...

class SumDigitsTask(Task):
    class Config(Task.Config):
        rubric = SumDigitsRubric.Config()
        renderer_env_config = RendererEnvConfig()    # loose, on the subclass
    def make_envs(self, *, example, group_size, renderer):
        return [RendererEnv(message_env=SumDigitsEnv(env_input=example.env_input),
                            renderer=renderer, config=self.renderer_env_config)
                for _ in range(group_size)]

After — one task, the base is concrete, the env is part of the config:

# RLTrainer.Config
task = SumDigitsTask.Config()
group_size = 8                                        # explicit (was overloaded onto sampling.n, see §4)
# TODO: support multiple tasks for data mixing.

class Task(Configurable):                             # concrete, make_envs is generic
    class Config(Configurable.Config):
        train_dataset: Configurable.Config
        validation_dataset: Configurable.Config
        rubric: Rubric.Config
        message_env: MessageEnv.Config               # the env is part of the Task config
        env_wrapper_cfg: RendererWrapperEnv.Config = field(default_factory=RendererWrapperEnv.Config)
    def sample_train_example(self):
        return next(self._train_dataset)             # datasets are iterators
    def sample_validation_example(self):
        return next(self._validation_dataset)
    def make_envs(self, *, example, group_size, renderer):
        return [RendererWrapperEnv(message_env=self._message_env_config.build(env_input=example),
                                   renderer=renderer, config=self._env_wrapper_cfg)
                for _ in range(group_size)]

# override methods to customize the Task beyond the config
# e.g. manage a pool or envs instead of making them fresh
class SumDigitsTask(Task):
    class Config(Task.Config):
        train_dataset = SumDigitsDataset.Config(seed=42)
        validation_dataset = SumDigitsDataset.Config(seed=99)
        rubric = Rubric.Config(reward_fns=[RewardCorrect.Config(weight=1.0), RewardFormat.Config(weight=0.3)])
        message_env = SumDigitsEnv.Config()

3. Reward functions are Configurable and live in config

"can they just define a Rubric and use it in a config, instead of register" · "can we just put them in config" (weights) · "any relationship between reward and components… can reward be a cached_property of components?"

Change: A RewardFn is one configurable scoring criterion. A Rubric is the concrete scheme that builds and weights them. The common case needs no subclass and is CLI-tunable.

Before — subclass + register_funcs(), weights hardcoded:

class Rubric(Configurable, abc.ABC):
    @abc.abstractmethod
    def register_funcs(self) -> list[RewardFn]: ...           # RewardFn = dataclass(fn=…, weight=…)

class SumDigitsRubric(Rubric):
    def register_funcs(self):
        return [RewardFn(fn=reward_correct, weight=1.0), RewardFn(fn=reward_format, weight=0.3)]

After — RewardFn is a Configurable callable, Rubric is concrete, the rubric is just config:

class RewardFn(Configurable, abc.ABC):
    class Config(Configurable.Config):
        weight: float = 1.0                                   # + any args a stateful reward needs
    async def __call__(self, rollout, env_input) -> float: ...

class Rubric(Configurable):
    class Config(Configurable.Config):
        reward_fns: list[RewardFn.Config] = field(default_factory=list)
        truncation_reward: float | None = None
        error_reward: float | None = None

A concrete reward function — the whole thing is the __call__ and a weight from config:

class RewardCorrect(RewardFn):
    """1.0 if the last [ANSWER] <n> equals the target, else 0.0."""
    class Config(RewardFn.Config):
        pass                                                  # only needs `weight`
    async def __call__(self, rollout, env_input: SumDigitsExample) -> float:
        matches = _ANSWER_RE.findall(last_assistant_text(rollout))
        if not matches:
            return 0.0
        return 1.0 if int(matches[-1]) == env_input.target else 0.0

# config:
Rubric.Config(reward_fns=[RewardCorrect.Config(weight=1.0), RewardFormat.Config(weight=0.3)])

On reward vs reward_breakdown — can reward be a cached_property of the breakdown? No: the rubric decides reward, and a custom score_group can set it however it likes (pairwise, rank-normalized, an LLM judge), so reward is not a function of the breakdown. reward_breakdown is the raw per-fn values, kept as its own field because callers consume it directly (e.g. per-component advantage), not just for logging. By default the Rubric happens to set reward to a normalized weighted sum of those values, but that is only the default.

4. Three small config knobs: group size, renderer, thread pool

"can we not let user config n and always fix it to be 1?"

done

· "Not clear what this means… avoid transformers dependency"

i rewrote the renderer config. Hopefully is clear now

· "This really sounds details that users shouldn't need to care" (the executor pool)

I removed from config and used a default with an educated guess based on number of rollouts

5. Completions map back by a stable id, and prompt-too-long is explicit

"comment is rollout_index, but code uses prompt_idx. What's the definition of prompt_idx?" · "not clear what 'overflow' means here"

Change: Each prompt is tagged with its rollout's sample_id.

Before

Completion.prompt_idx: int
Rollout.sample_idx: int
ordered = sorted(completions, key=lambda c: c.prompt_idx)        # brittle int sort

After

sample_id = f"{group_id}/sample={sample_idx}"                    # readable id at the boundary
Completion.request_id: str
Rollout.sample_id: str
returned_ids = [c.request_id for c in completions]
if returned_ids != request_ids:                                 # reorder / missing / n>1
    raise RuntimeError(...)

RolloutStatus.TRUNCATED_PROMPT_OVERFLOW → TRUNCATED_PROMPT_TOO_LONG (and _is_prompt_overflow → _is_prompt_too_long).

[felipemello1]

@tianyu-l

the new layout

torchtitan/experiments/rl/
├── environment/        # message space ⇄ token space   (was env_types/)
│   ├── message.py      #   MessageEnv  (+ MessageInitOutput / MessageStepOutput)
│   └── token.py        #   TokenEnv    (+ TokenEnvOutput)
├── rollout/            # rollout data + the rollout producer   (was rollouts/)
│   ├── types.py        #   Rollout / RolloutGroup / RolloutTurn / RolloutStatus
│   ├── utils.py
│   └── rollouter.py    #   Rollouter           (was tasks/task.py → Task)
├── examples/           # concrete problems      (was tasks/)
│   └── sum_digits/     #   data.py / env.py / rubric.py / rollouter.py (SumDigitsRollouter)
├── actors/
│   ├── generator.py    #   VLLMGenerator       (TODO: → InferenceEngine)
│   └── trainer.py      #   PolicyTrainer
├── rubrics/rubric.py   #   Rubric / RubricOutput
├── grpo.py             #   RLTrainer + GRPOLoss (controller / train loop)
└── renderer.py  batcher.py  generate.py  config_registry.py  types.py

Goal for the code (multi-turn, once the generator supports it):

async def do_single_rollout(generator, env) -> Rollout:
    turns = []
    env_output = await env.init()
    while not env_output.status.is_terminal():
        prompt_token_ids = env_output.next_prompt_token_ids
        prompt_messages = env_output.next_prompt_messages
        completion = await generator.generate(prompt_token_ids)
        env_output = await env.step(completion)
        turns.append(
            RolloutTurn(
                prompt_token_ids=prompt_token_ids,
                prompt_messages=prompt_messages,
                completion_token_ids=completion.token_ids,         # generator  -> completion_*
                completion_logprobs=completion.token_logprobs,
                parsed_completion_message=env_output.parsed_completion_message,
                env_messages=env_output.env_messages,
                env_rewards=env_output.env_rewards,
            )
        )
    return Rollout(turns=turns, status=env_output.status)

---

1. environment/ package; the token wrapper is TokenEnv

"prefer full spelling… rl/environments/message.py" · "if it's doing more than rendering, the additional part would justify a broader name" · "if you follow MessageEnv, this should be called TokenEnv"

• env_types/ → environment/ (message.py + token.py) (note: called it "environment" singular, to not confuse with a place to put environments like sumdigits)
• RendererWrapperEnv → TokenEnv (renders + parses + owns the rollout checks; now a Configurable)
• TokenizedStepOutput → TokenEnvOutput (not RawOutput/RenderedOutput — Message*/Token* names the space, which is the actual distinction)

2. init not reset; message vs prompt

"reset sounds like cleaning up an old state… maybe init" · "reserving prompt for the first message is ambiguous"

• reset → init (these envs are single-use)
• MessageResetOutput.prompt_messages → MessageInitOutput.init_prompt_messages (qualified with init so it's unambiguously the first-prompt messages)

3. the generator's output is completion

"assistant vs generator used interchangeably?" · "does completion.token_ids contain the full history?"

Everything the generator makes is completion_*:

• assistant_token_ids / assistant_logprobs → completion_token_ids / completion_logprobs
• assistant_message → parsed_completion_message

4. SamplingConfig is user knobs only

"can we not make this configurable" (n) · "I'd rather not put stop_token_ids in SamplingParams, send it from the engine constructor"

done — dropped both:

• SamplingConfig.n → removed (n=1 is hardwired; group_size pre-expands the prompts)
• SamplingConfig.stop_token_ids → removed (injected into the generator's constructor instead)
• also killed the extra replace() in RLTrainer.__init__

5. Task → Rollouter

"is it too bad if we call this RolloutGenerator, mimicking Dataloader in SFT… I'm even willing to rename Generator to InferenceEngine for this."

• renamed Task → Rollouter
• also SumDigitsTask → SumDigitsRollouter
  -RLTrainer.Config.task → rollouter.
• package moves (see the layout up top): tasks/ dissolved — the base Rollouter now lives in rollout/rollouter.py, concrete problems moved to examples/sum_digits/;
• VLLMGenerator → InferenceEngine left as a # TODO(naming) at the class — the "rename Generator" you offered (broad cross-tree rename; separate PR).

6. rollouter construction

"make_env_group, as you have score_group" · "env is ambiguous"

• make_envs → make_env_group
• env_wrapper_cfg → token_env

7. the nits

• _terminal deleted (inlined)
• Rollout.status → required (no COMPLETED default)
• num_prompts_per_step → num_groups_per_rollout_batch
• gen_metrics → generation_metrics
• added a comment on how max_concurrent_rollouts is derived (and what concurrency it bounds)
• TokenEnvOutput locals → env_output / env_init_output
• .get() → await on the rollout generate call
• trimmed the unnecessary list() copies
• truncation_reward/error_reward docstrings clarified
• shape hints spelled out: [L_prompt]/[M_prompt]/[L_response] → [num_prompt_tokens]/[num_prompt_messages]/[num_completion_tokens] (also [num_env_messages]/[num_tools]/[num_turns])

8. kept on purpose (pushed back, with reasoning)

"prompt seems a term in the resource space; message is more of a user-facing thing -- can we not mix up the two?"

I decided to keep the prompt, since the MessageEnv is aware of the concept. For "init" it produces a prompt. For "step", it has to produce the next prompt. Let me know if you strongly disagree.

"this check should be added for init output as well?"

No — MessageStepOutput rejects assistant turns because a step reply is tool/user only. The init conversation can legitimately carry few-shot assistant turns.

"storing turns: list[RolloutTurn] in Rollout sounds quite redundant, O(n) -> O(n^2)."

Full-history-per-turn is a common layout (verifiers, rllm, tinker). We collapse these at Episode assembly. Keeping the full history is helpful to find divergences and branch (e.g. for compactation). If memory becomes an issue, then we can revisit. Let me know if you disagree.

"should they be put into a Reward"

we could do something like:
Rollout.reward = RubricOutput(reward, reward_breakdown)

I opted to keep it flat for now. Let me know if you strongly disagree.

"I'd rather we create a thin Renderer class inheriting prime-rl's renderer, and have this as Renderer.Config."

There is nothing really to subclass. The 'config_from_name' returns a model specific class, e.g. Qwen3Rendererconfig.
The correct API here would be to do:

renderer = create_renderer(tokenizer, Qwen3RendererConfig(enable_thinking=False))

We do:

class RendererConfig(Configurable.Config):
    def build(self, *, tokenizer_path: str) -> Renderer:
        renderer_config: Qwen3RendererConfig = config_from_name(renderer_name)
        ... # override params
        renderer = create_renderer(tokenizer, renderer_config_with_overrides)

"Can we do this (Thread pool) in Renderer's constructor, after we create one."

I don't think that loop-wide concurrency should be hidden in the renderer. Its better to be exposed in the controller, no? Let me know if you strongly disagree

follow-ups (TODOs in code, separate PRs)

• generator_logprobs → old_logprobs at the loss (vs policy_logprobs)
• VLLMGenerator → InferenceEngine (TODO left at the class)

[wwwjn]

Thanks for the comments, it's much more clear

Goal for the code (multi-turn, once the generator supports it)

I might missed something, what's the gap for generator to support this "code goal"?

[felipemello1]

Thanks for the comments, it's much more clear

Goal for the code (multi-turn, once the generator supports it)

I might missed something, what's the gap for generator to support this "code goal"?

allow multiple threads to call await generator.generate{prompt_request) independently.

Currently we have to batch all requests because our generator doesn't have any await that releases the lock for other threads to add more requests to the queue.

I will have a PR to fix that probably tomorrow

[wwwjn]

Why we need this type checking?

[wwwjn]

What does this "running conversation" mean?

[wwwjn]

This means the model generate a <tool_use> token, right?

and I want to do a knowledge check of my mental model. say if we are training on a Coding task:

• Do we need to plug both SandBoxEnv and MessageEnv into the TokenEnv
• When there's a <tool_use>` token, TokenEnv will call the SandBoxEnv to execute the tool?

I guess my questions is: where will the tool calling happen in the future?

[wwwjn]

Chatted with Claude, it says:

 In this design, MessageEnv is the problem-specific environment that receives the
  assistant’s parsed message and decides what happens next. For a coding task, that env
  would own or call the sandbox.

  VLLMGenerator
    |
    | completion tokens
    v
  TokenEnv
    |
    | Renderer.parse_response(...)
    | parsed assistant message
    v
  CodingMessageEnv.step(parsed_completion_message)
    |
    | extract code / tool call / patch / command
    v
  Coding Sandbox
    |
    | run tests, execute code, inspect files, enforce timeout
    v
  CodingMessageEnv
    |
    | returns tool/user messages + optional env_rewards + done
    v
  TokenEnv
    |
    | render next prompt tokens if not done
    v
  VLLMGenerator

I updated my mental model: TokenEnv is plumbing between Generator and message space

[tianyu-l]

I promise this is last push.

[tianyu-l]

I think after we name Rollouter, this is less necessary.

[tianyu-l]

Since you have TokenEnvOutput, I would mirror this to MessageEnvInitOutput, especially considering it contains not just messages but also tools.

[tianyu-l]

mention that we don't post_init check because assistant is legit

[tianyu-l]

Is this part of next_prompt_messages? like

next_prompt_messages contains
[previous prompt_messages, parsed_completion_message, env_messages]

[tianyu-l]

can we use completion_tokens vs completion_message?

[felipemello1]

the reason i added "parsed" is to make it evident that it did NOT come from the vllm completion, but rather it was parsed by the tokenEnv. Does this make sense? I am ok with renaming, but we lose this detail

[tianyu-l]

That's fair, but for perfect consistency it sounds we should do

• prompt_message
• tokenized_prompt_token_ids
• completion_token_ids
• detokenized_completion_message

There are a lot of redundancy so I'm not sure.

I would like to at least avoid inventing another word "parsed"

[tianyu-l]

Looks I'm mostly correct.

[felipemello1]

now that i am thinking about it, this is at the wrong place. The MessageEnv should say what the "next_messages" are, e.g. for compacting history. Let me see how i can move this logic to MessageEnv. I guess this would mean i have to delete "env_messages" and just have it return "next_prompt_messages"

[tianyu-l]

It seems you could've put these in TokenEnvOutput, so that we can consolidate RolloutTurn and TokenEnvOutput.

If so, I would strongly advocate that, in this version.

[felipemello1]

class TokenEnvOutput:
    # off by one
    next_prompt_token_ids: list[int] | None  # [num_prompt_tokens] or None
    next_prompt_messages: list[Message] | None = None  # [num_prompt_messages] or None

	# different 
    status: RolloutStatus

	# Same
    parsed_completion_message: Message | None = None
    env_messages: list[Message] = field(default_factory=list)  # [num_env_messages]
    env_rewards: dict[str, float] = field(default_factory=dict)

class RolloutTurn:
	# off by one
    prompt_token_ids: list[int]  # [num_prompt_tokens]
    prompt_messages: list[Message] = field(default_factory=list)  # [num_prompt_messages]

	# different 
    policy_version: int | None = None
    completion_token_ids: list[int]  # [num_completion_tokens]
   	completion_logprobs: list[float]  # [num_completion_tokens]

	# Same
    parsed_completion_message: Message | None = None
    env_messages: list[Message] = field(default_factory=list)  # [num_env_messages]
    env_rewards: dict[str, float] = field(default_factory=dict)

RolloutTurn is aware of the whole rollout loop. There may be other generators (e.g. a teacher), more metrics, some crazy logic. While the TokenEnvOutput is concerned only about the Env.

I think it would be too opinionated to do consolidate them and would put unnecessary responsability on the TokenEnv.

I do agree that there is some overlap. We could create some dataclasses to hold some of the fields that are common, e.g. Completion or tokens_with_logprobs can be an output of the generator and we do RolloutTurn.completion.

NextPrompt could hold next_prompt_token_ids and next_prompt_messages.

cc: @wwwjn in case you have an opinion as well

[tianyu-l]

sounds good to leave them separate right now

[tianyu-l]

not strongly opinionated

Suggested change

	example: object,
	sample: object,

[tianyu-l]

nit: use "training" instead of "train" to be consistent

Suggested change

	def sample_train_example(self) -> object:
	def get_training_sample(self) -> object:

[tianyu-l]

prefer to not have this indirection.

[tianyu-l]

don't need to pass this in, you have self._rollouter

[felipemello1]

for datamix we would need it, but all of this code in grpo will go under the rollout anyway, and in this case, it would just use self. Either way, for now, we have 1 rollout, so i agree

[wwwjn]

A side note: I read from somewhere else the request_id needs to contain the DP rank id, otherwise each DP rank will have the same set of request_ids, which will causing collision request id error in vllm engine. I think it's a good practice to make the request id globally unique.

We can put a TODO here, wdty?

[felipemello1]

I am not sure if i follow: The id is unique. I guess i could add some random letters to make sure.

But would we send the same sample_id to different dp ranks? Shouldnt each rank get a different set of ramples?

[tianyu-l]

same prompt go to same DP rank for prefix cache reuse

[wwwjn]

Is completions sorted?

Sort and then slicing to group samples into samples are error-prone, can we get group_id from request_id and then group using group id?

[felipemello1]

They are sorted. I should fix it. But since this code will be delete/refactored in the next feel days, is it ok if i leave a todo for now?

[wwwjn]

What's the consideration of putting _run_group_rollout, _run_single_rollout in controller vs. in Rollouter?

if we put it in rollouter, I would image controller could simply call

self._rollouter.run_group_rollout(generator, sample_id, xxx)

[felipemello1]

yes, thats the direction. I thought about already removing the "collect_rollouts" from grpo and put it in the Rollouter. I was going to leave it for the next PR refactor. Maybe we should land as is to avoid another round of review that will be wasted anyway due to the refactor?

Now i feel like i should have fixed the generator first before doing this PR haha

[felipemello1]

All requests addressed

Only one that I didnt was MessageEnv returning "next_prompt_messages" for history-editing cases, like compacting. This is not trivial. We would need a postprocessing function in rollout_to_episode to branch messages if history diverges, which I will add in the multiturn PR

Also, it might be easier in the single_rollout_loop to intercept TokenEnvOutput, run the compaction service and create a new TokenEnvOutput.

For now, i added a TODO

cc: @tianyu-l @wwwjn

[tianyu-l]

Great work!!

---

Not sure if you missed #3453 (comment)

The benefit is that

• user knows what they should change (env, rubric, etc.), what they should not (generator / trainer)
• consistent with pretraining https://github.com/pytorch/torchtitan/tree/main/torchtitan/protocols

Shouldn't block this PR. Happy to discuss more.