Training Agentic & Reasoning Models

How to fine-tune models that think, use tools, and recover from errors — the data format, the masking decisions, and the ECHO observation loss.

---

What makes agentic data different

Standard SFT data is simple: user asks, assistant answers. The loss is on the answer.

Agentic data has structure:

user → assistant thinks → calls tool → gets result → thinks again → calls another tool → ... → final answer

Each turn has different semantics. The assistant's thinking is something you want the model to learn. The tool calls are actions you want it to reproduce. The tool responses are environment observations that the model didn't generate — so should they get loss or not?

This question — what gets loss? — determines what the model learns.

---

Data format

Palingenesis supports the full agentic trace structure:

{"messages": [
  {
    "role": "system",
    "content": "You are a coding assistant with access to file editing tools."
  },
  {
    "role": "user",
    "content": "Fix the bug in the group_by function."
  },
  {
    "role": "assistant",
    "content": "",
    "reasoning": "Let me start by exploring the repository structure to understand the codebase. The issue describes that group_by only groups nearby elements, and we need a new chunk_by method.",
    "tool_calls": [{"function": {"name": "view_file", "arguments": "{\"path\": \"/src/lib.rs\"}"}}]
  },
  {
    "role": "tool",
    "content": "pub fn group_by<K, F>(self, key: F) -> GroupBy<K, F, Self> ..."
  },
  {
    "role": "assistant",
    "content": "I've found the issue. The current implementation uses...",
    "reasoning": "Now I can see the problem. The GroupBy struct only tracks the current key, not previous keys. I need to create a new ChunkBy struct that maintains a map of seen keys."
  }
]}

Tool definitions go in a tools field next to messages (a list, or a JSON string: Arrow often cannot store heterogeneous JSON schemas natively). They are passed to the chat template as tools=, which renders them into the system prompt the way the model sees them at inference. The field name is data.tools_field (default tools).

Tool-call arguments may be a dict or a JSON string (the OpenAI wire format); strings are parsed, since chat templates iterate over the arguments as a mapping. A whole conversation stored as a JSON string also works.

Per-message training flags

A message with "loss": false (or "weight": 0) stays in the context but receives no loss, e.g. a canned greeting, or a failed attempt kept so the model learns the recovery that follows:

{"role": "assistant", "content": "Hi! How can I help you today?", "loss": false}

Messages without the flag (or with "loss": null, which Arrow fills in for messages that lack the key) are trained.

Role semantics and masking

Role	Field	Gets loss?	Why
system	content	No	Context the model doesn't generate
user	content	No	Input the model doesn't generate
assistant	content	Yes	The visible output — must be reproduced
assistant	reasoning	Yes	The thinking — model learns to reason (legacy reasoning_content also read)
assistant	tool_calls	Yes	The actions — model learns when/how to use tools
tool	content	No (default)	Environment response, not generated by model
tool	content	Yes (ECHO mode)	See below

The chat template handles the serialization. reasoning is passed to the template under both reasoning and reasoning_content (as vLLM does), so Qwen3/3.5 templates — which read reasoning_content — render it as <think>...</think> tokens. For other models, it's part of the assistant turn.

An assistant turn is trained as the template renders it: everything after the template's assistant header, up to and including the end-of-turn token. That covers the tool-call markup (<tool_call>...</tool_call> on Qwen), which templates render from tool_calls rather than from content, and the end-of-turn token after it, which teaches the model to stop and wait for the tool result. An empty <think></think> block that a template inserts into a turn without reasoning is not trained.

---

Standard mode (default)

data:
  include_observations: false

Only assistant tokens (content + reasoning + tool calls) receive loss. Tool responses are context — they inform the next assistant turn but the model isn't trained to predict them.

This is correct for most scenarios. The model learns: - How to reason (from reasoning) - When to call tools (from tool_calls) - How to use results (from seeing tool responses as context, then producing the next action)

But it does NOT learn: what tools return. It's a pure action model.

---

ECHO mode: observation loss

data:
  include_observations: true

From "Terminal Agents Learn World Models for Free" (ICML 2026, arxiv:2605.24517):

When you also train on tool responses, something unexpected happens: the model develops an internal world model. It learns to predict what tools will return before actually calling them. This has two concrete benefits:

1. Better tool selection: the model can "simulate" a tool call internally and decide if it's worth executing
2. Error recovery: when a tool returns unexpected output, the model has calibrated expectations and can reason about the discrepancy

The paper shows this emerges naturally — you don't need to engineer it. Just adding tool response tokens to the loss is sufficient.

When to use ECHO

• Yes: training coding agents (file contents are predictable), search agents (results follow patterns), API agents (schemas are learnable)
• No: training on traces where tool responses are unpredictable noise (e.g., random UUIDs, timestamps, non-deterministic API responses)

How it interacts with DEFT

DEFT (the adaptive token-weighting loss) composes perfectly with ECHO. The model will naturally assign low weight to unpredictable observation tokens (high perplexity → DEFT downweights) and high weight to predictable ones (file structures, known API patterns). You get the world-model benefit without manually filtering observation quality.

data:
  include_observations: true
plugins:
  deft: true  # Automatically downweights noise in observations

---

Turn scaling

Multi-turn agentic traces often have 5-20 turns. Not all turns are equally important. The final turn (where the agent commits to a solution) matters more than intermediate exploration.

data:
  turn_scaling: progressive   # Later turns get more weight

Options:

Mode	Effect	Best for
uniform	All turns equal weight	Short conversations (1-3 turns)
progressive	Weight = √(turn_idx / total_turns)	Agentic traces (iterative refinement)
last_heavy	Final turn gets 2× weight, others 1×	Task completion (answer quality matters most)

The default is uniform. The case for progressive: in a 10-turn trace, the early turns are often exploration (trying things, getting errors) and the later turns produce the real solution. This is a heuristic, not a result from a paper; measure it on your eval set. Weights are normalised to mean 1 per conversation, so the loss scale is unchanged. They are applied by the cross-entropy losses (plain, chunked, Cut Cross-Entropy) and by chunked DEFT; other objectives, SeCO and DPO reject a non-uniform turn_scaling.

---

Example config for agentic traces

agentic config (Qwen3.5, long traces)

model:
  name_or_path: Qwen/Qwen3.5-4B
  torch_dtype: float32           # SeCO's recommendation; compute runs in bf16
  compile: false                 # SeCO is validated without torch.compile

data:
  dataset: ./agentic_traces.parquet   # messages (+ tools) per row
  streaming: false
  max_seq_length: 65536          # whole traces; SeCO keeps memory set by the chunk
  include_observations: false    # set true for ECHO (trains on tool outputs)

train:
  epochs: 1
  per_device_batch_size: 1
  gradient_accumulation_steps: 16
  learning_rate: 1.0e-5
  gradient_checkpointing: full

memory:
  seco: true
  seco_chunk_size: 4096

For traces that fit in memory without SeCO (a few thousand tokens), drop the memory: section and use gradient_checkpointing: selective; the chat-template masking, tools and per-message flags work the same way.

---

Debugging masking

Before training on agentic data, always verify that masking is correct:

pgs inspect --config configs/agentic_traces.yaml --num_samples 3

This shows you exactly which tokens have loss (highlighted) and which are masked. For agentic traces, you should see:

• System prompt (including the tool definitions): masked
• User message: masked
• Assistant thinking (<think>...</think>): has loss
• Assistant tool calls, and the end-of-turn token after them: has loss
• Tool responses: masked (or has loss if ECHO enabled)
• Assistant final answer: has loss
• Turns flagged "loss": false: masked

If anything looks wrong (e.g., user tokens getting loss, or thinking tokens masked), the chat template isn't matching your data format. Check your messages_field and ensure roles are spelled correctly.

---

How much reasoning data do you need?

From the research on reasoning-trace SFT:

• 500-2,000 traces: the model learns the format (when to think, how to structure tool calls) but doesn't develop deep reasoning
• 5,000-20,000 traces: genuine reasoning capability emerges. The model starts transferring problem-solving patterns across tasks.
• 50,000+ traces: diminishing returns unless traces are highly diverse across tools and domains

Quality matters enormously. One well-annotated trace with clear reasoning > 10 traces where the assistant stumbles to an answer by luck.

Use pgs prepare --strategy optimal to select the highest-value traces from a larger pool.

---

Data validation

Before training, validate that your tool calls are consistent with your declared tools:

python -c "
from palingenesis.validate_data import validate_dataset
import json

with open('data.jsonl') as f:
    samples = [json.loads(l) for l in f]

report = validate_dataset(samples, messages_field='messages', tools_field='tools')
print(f'Valid: {report[\"valid\"]}/{report[\"total\"]}')
print(f'Errors: {len(report[\"errors\"])}')
for idx, err in report['errors'][:10]:
    print(f'  Sample {idx}: {err}')
if report['tools_never_called']:
    print(f'Tools declared but never called: {report[\"tools_never_called\"]}')
"

The validator checks:

• Every tool_calls[].function.name exists in the sample's tools list
• Every parameter passed to a tool exists in that tool's schema
• All required parameters are present in the call

Catching these errors before training prevents the model from learning to call non-existent tools or pass wrong arguments.

---

Role mapping (non-standard formats)

Your data doesn't need to use standard OpenAI roles. Palingenesis auto-normalizes during training:

Your format	Normalized to
{"from": "human", "value": "..."}	{"role": "user", "content": "..."}
{"from": "gpt", "value": "..."}	{"role": "assistant", "content": "..."}
{"role": "bot", "content": "..."}	{"role": "assistant", "content": "..."}
{"role": "function", "content": "..."}	{"role": "tool", "content": "..."}
Field "think" on assistant	Mapped to reasoning
Field "function_call" on assistant	Mapped to tool_calls
Conversation in "conversations" field	Found automatically

This means datasets in ShareGPT format, Alpaca format, or any other common layout work without conversion. The normalization happens transparently in the data pipeline.

If you have a truly custom mapping, you can extend it:

from palingenesis.validate_data import normalize_messages, ROLE_MAP

# Add your custom role names
ROLE_MAP["instructor"] = "user"
ROLE_MAP["learner_response"] = "assistant"