[Warning] When assemble steps, detect the trajectory not accumulative

[kxfan2002]

Thanks for your great work.

I'm running the math_tool training example. I first encounter this:

File "envs/rllm2/lib/python3.11/asyncio/tasks.py", line 489, in wait_for
  return fut.result()
     ^^^^^^^^^^^^
 File "RL/rllm/rllm/engine/agent_execution_engine.py", line 394, in run_agent_trajectory_async
   prompt_tokens, response_tokens, response_masks, is_valid_trajectory = self.assemble_steps(episode_steps)                                                                        ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "RL/rllm/rllm/engine/agent_execution_engine.py", line 453, in assemble_steps
   current_completion_ids = step["completion_ids"]
                        ~~~~^^^^^^^^^^^^^^^^^^
 KeyError: 'completion_ids'

I changed line253 into "completion_ids": model_output.completion_ids, in rllm.engine.agent_execution_engine. I guess this is a little mistake that should be correct in the repo.

Then I see this warning during training:

(TaskRunner pid=32375) WARNING:2025-12-17 21:07:06,390:When assemble steps, detect the trajectory not accumulative at position 4688. Expected: [1699, 25310, 367, 284, 33122], Got: [59, 811, 446, 367, 284]. Setting response_masks to all 0s. This is likely due to retokenization.
(TaskRunner pid=32375) Trajectory 50 completed due to: ENV_DONE. Reward is 0.0. 
(TaskRunner pid=32375) 
(TaskRunner pid=32375) Number of Trajectories 199/256 completed
(TaskRunner pid=32375) Trajectory 53 completed due to: ENV_DONE. Reward is 0.0. 
(TaskRunner pid=32375) 
(TaskRunner pid=32375) Number of Trajectories 200/256 completed
(TaskRunner pid=32375) WARNING:2025-12-17 21:07:10,227:When assemble steps, detect the trajectory not accumulative at position 7240. Expected: [1699, 1830, 284, 220, 15], Got: [59, 1016, 629, 284, 220]. Setting response_masks to all 0s. This is likely due to retokenization.
(TaskRunner pid=32375) Trajectory 55 completed due to: ENV_DONE. Reward is 0.0. 
(TaskRunner pid=32375) 
(TaskRunner pid=32375) Number of Trajectories 201/256 completed
(TaskRunner pid=32375) Trajectory 49 completed due to: MAX_STEPS. Reward is 0. 
(TaskRunner pid=32375) 
(TaskRunner pid=32375) Number of Trajectories 202/256 completed
(TaskRunner pid=32375) WARNING:2025-12-17 21:07:16,377:When assemble steps, detect the trajectory not accumulative at position 6668. Expected: [1699, 5035, 284, 220, 15], Got: [59, 406, 2370, 284, 220]. Setting response_masks to all 0s. This is likely due to retokenization.

I wonder why this happens. Is this an implementation error or a tool calling caused issue?

When assemble steps, detect the trajectory not accumulative at position 4688. Expected: [1699, 25310, 367, 284, 33122], Got: [59, 811, 446, 367, 284]. Setting response_masks to all 0s. This is likely due to retokenization.

I would appreciate it if you could reply. Thanks again for your great contribution!

[listar2000]

This is crucial and I've also encountered this issue (seems to be introduced due to the recent v0.2.1). Do you want to send a PR (I think your one-line fix is good)?

For the second issue, I will take a look at this.

[jeffreysijuntan]

the latest v0.2.1.post release should address this.

[kxfan2002]

the latest v0.2.1.post release should address this.

Hi and thanks for the quick reply!

I see the commit history. The latest v0.2.1.post changed code about local search server in example/search/retrieval. However, I encountered

When assemble steps, detect the trajectory not accumulative at position 4688. Expected: [1699, 25310, 367, 284, 33122], Got: [59, 811, 446, 367, 284]. Setting response_masks to all 0s. This is likely due to retokenization.

in example/math_tool too. I think the latest release is not effective for this? If I miss anything, please tell me.

Thanks again for your contribution and quick reply!

[kxfan2002]

I think this is because some procedure in training somehow "rewrite" the trajectory tokens. If you have an app called xiaohongshu, I think this post(💥agentic RL务必记得token-in-token-out 最近做 Agen... http://xhslink.com/o/7YLGfVURj8g
复制后打开【小红书】查看笔记！) maybe relevent (?)

[kxfan2002]

I print out the mismatch tokens in the warning:

exp=[1699, 25310, 367, 284, 33122]
got=[59, 811, 446, 367, 284]

exp=[1699, 1830, 284, 220, 15]
got=[59, 1016, 629, 284, 220]

exp=[1699, 5035, 284, 220, 15]
got=[59, 406, 2370, 284, 220]

The results are:

\nequation = Eq
\nequation =

\ncount = 0
\ncount = 

\ntotal = 0
\ntotal = 

It seems like this is part of the code model writes in python tool use.

[jeffreysijuntan]

Got it, investigating this issue now. We switched to token-in-token-out mode in v0.2 in our workflow engine, but there could be some compatibility issue with the prior AgentExecutionEngine that is causing this behavior.

[jeffreysijuntan]

Since the math_tool example are still using the AgentExecutionEngine, not the AgentWorkflowEngine, retokenization still happens because the agent returns a Action class that contains a string we retokenize over. That's why we add this check to detect discrepancy in the trajectory. If we detect a mismatch after retokenization, this trajectory will not be trained. This is not a perfect solution, but should prevent retokenization drift from happening.

We are currently working to unify the AgentExecutionEngine and the AgentWorkflowEngine. In the future, both will use AgentWorkflowEngine as the backbone so this issue won't happen any more.

[XiaofengLin7]

rllm/rllm/engine/agent_workflow_engine.py

Lines 274 to 285 in 186d0bb

	if not self.config.rllm.stepwise_advantage.enable:
	if len(trajectory.steps) > 1:
	if not trajectory.is_cumulative():
	logger.warning(f"Warning: Multi-step trajectory {trajectory_id} is not cumulative, but stepwise mode is not enabled. There could be a token mismatch during trajectory generation.")
	chat_completions = trajectory.steps[-1].chat_completions
	chat_completions_list.append(chat_completions)
	prompt, response, mask = self.rollout_engine.chat_parser.tokenize_and_mask_cumulative(chat_completions)
	prompts.append(prompt)
	responses.append(response)
	traj_mask.append(mask)
	multi_modal_inputs_list.append({}) # empty dict

Retokenization still happens if we set self.config.rllm.stepwise_advantage.enable as False and len(trajectory.steps) > 1? Is this intended?