huggingface / huggingface/trl

# [Feature Request] GRPOTrainer: support stop-tool termination in agent training loop

Open
#5,389 1 comment 0 reactions 0 assignees View on GitHub
Dominant language
Python
Stars
19.3k
Forks
3k
Avg merge
1d 20h
Merged PRs (30d)
194

Description

# [Feature Request] GRPOTrainer: support stop-tool termination in agent training loop

## Summary

When training agents with `GRPOTrainer` using `tools` + `max_tool_calling_iterations`, there is no way to designate a tool as a **termination signal** — a tool whose invocation means "the agent is done, do not generate any further". This creates a systematic divergence between training and evaluation when using frameworks like [smolagents](https://github.com/huggingface/smolagents) that rely on a `final_answer` tool to end agent runs.

## Current behaviour

The agent loop in `_tool_call_loop` terminates only when:
1. `max_tool_calling_iterations` is reached, or
2. The model generates a turn with no tool calls (EOS).

If the model calls `final_answer(answer=42)`, TRL executes it, appends the result to the context, and then calls `_generate_single_turn` again — the model must now generate *another* turn before the loop can exit. This wastes context budget (potentially hitting `max_completion_length` before EOS), and more importantly, it means the agent's training behaviour differs from its evaluation behaviour.

## Proposed behaviour

Add a `stop_tool_names: list[str] | None` field to `GRPOConfig`. When a sample's tool call matches any name in this list:

1. The tool is executed normally and its result is appended to the context.
2. `_generate_single_turn` is **not** called for that sample — it exits the loop immediately.
3. The `tool_mask` and `completion_ids` are updated to reflect the tool result tokens (masked out of loss), exactly as they would be for a normal tool call.

This mirrors the `final_answer` semantics used in smolagents and similar agentic frameworks.

## Proposed API

```python
from trl import GRPOConfig, GRPOTrainer

def final_answer(answer: int) -> int:
"""Submit the final answer."""
return answer

config = GRPOConfig(
...,
stop_tool_names=["final_answer"],
max_tool_calling_iterations=10,
)

trainer = GRPOTrainer(
...,
args=config,
tools=[tool_a, tool_b, final_answer],
)
```

With this config, as soon as any sample's model calls `final_answer`, that sample exits the agent loop. Other samples in the batch continue their own tool-calling turns as usual.

## Motivation: training/eval alignment

When evaluating with smolagents, `final_answer` is a built-in termination tool. If the model is trained with TRL *without* stop-tool support, it learns to:
- Call `final_answer` → observe the result echoed back → generate more tool calls → eventually hit `max_tool_calling_iterations`

This is a different behaviour from what it will encounter at eval time (call `final_answer` → done). Over many gradient steps, the model's policy diverges from the target behaviour.

With `stop_tool_names`, training and evaluation share the same agent loop semantics.

## Implementation sketch

The change is contained in two files:

**`trl/trainer/grpo_config.py`** — add field:
```python
stop_tool_names: list[str] | None = field(
default=None,
metadata={"help": "Tool names that terminate the agent loop immediately when called."},
)
```

**`trl/trainer/grpo_trainer.py`** — in `_tool_call_loop`, after the tool execution block and after the overlong filter, before calling `_generate_single_turn`:
```python
# Remove samples that called a stop tool — finalize their state without further generation
if self.args.stop_tool_names and stop_tool_global_idxs:
non_stop_mask = [i not in stop_tool_global_idxs for i in idxs_with_tool]
for local_idx, idx_with_tool in enumerate(idxs_with_tool):
if idx_with_tool in stop_tool_global_idxs:
prompt_length = len(prompt_ids[idx_with_tool])
pct = prompt_completion_tool_ids[local_idx]
old_completion_length = len(completion_ids[idx_with_tool])
tool_suffix_length = len(pct) - prompt_length - old_completion_length
completion_ids[idx_with_tool] = pct[prompt_length:]
tool_mask[idx_with_tool] += [0] * tool_suffix_length
if logprobs is not None:
logprobs[idx_with_tool] += [0.0] * tool_suffix_length
idxs_with_tool = [i for i, keep in zip(idxs_with_tool, non_stop_mask) if keep]
prompt_completion_tools = [p for p, keep in zip(prompt_completion_tools, non_stop_mask) if keep]
prompt_completion_tool_ids = [p for p, keep in zip(prompt_completion_tool_ids, non_stop_mask) if keep]
if not idxs_with_tool:
break
```

I'm happy to submit a PR with tests if the approach looks good. Related: #2723, discussion #2704.

Contributor guide

Open the contributing guide

Assessment

This issue has not been assessed yet.

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.