# Lab 18: a tool-using agent, from rollout loop to masked loss

You will build the pieces of agentic RL that a single-response method like GRPO never needed:
an environment with a tool, a rollout loop that braids model tokens with tool tokens, credit
assignment across turns, and a loss that trains only on what the model wrote. Then you will
train a tiny agent and watch it learn to search with the tool, and watch it fail when you
forget the mask.

Chapter: <https://learn.xhelio.ai/rl/learn/18-agentic-rl/>

## The task

A secret number `s` in 0..7. The agent writes two-token turns:

- `call x` asks the tool `compare(x)`, which answers `lower` (s < x), `higher` (s > x) or `equal`. Cost: 0.05.
- `ans x` ends the episode: reward 1 if `x == s`, else 0.
- Anything else is malformed: the episode halts with reward −1.
- At most 5 turns.

Half the tasks show the number in the prompt (`[<bos> 5]`: the tool is useless). The other half
hide it (`[<bos> ?]`: the tool is the only way to do better than guessing). A trained agent
looks like this:

```text
[<bos> ?]  call 3  {<tool> higher}  call 5  {<tool> lower}  ans 4      turn rewards [-0.05, -0.05, 1.0]
 prompt    model   tool output      model   tool output     model
```

### Why not TinyLM?

Labs 11–17 use `tinylm.py`. Here a fresh 2-layer transformer would need new tool tokens, a format
it has never seen, and minutes of training per run, and the arithmetic task would mix "can the
model add?" with "does it use the tool?". Instead the policy is a tiny GRU language model (about
20k parameters, 17-token vocabulary) that trains in under a minute. Everything that matters for
agentic RL is the same as with a real LLM: the policy reads one flat token sequence, generates
token by token, conditions on tool output it did not write, and is trained on per-token
log-probabilities with a mask.

## Files

| File | What it is |
|---|---|
| `agentic.py` | **Your file.** Four tasks marked `TODO`. The training loop `train` at the bottom is complete. |
| `tinyagent.py` | Given: the vocabulary, the GRU `Policy` (trainer), a numpy `Sampler` (fast rollouts), `Trajectory`, the cold-start `warm_start`, `make_batch` (padding and masks), `render`. Read `make_batch` once. |
| `test_agentic.py` | Tests, grouped by task. |
| `run.py` | Trains three variants and saves `agentic_rl.png`. |
| `solutions/agentic.py` | Reference solution. Try not to peek. |

## Tasks

Run one task's tests at a time with `-k`, e.g. `python -m pytest -q -k task2`.

1. **`GuessEnv.step`** (Task 1). The environment: parse the turn, run the tool, return
   `(obs_tokens, reward, done)`. Tests check every rule, including the turn limit.
2. **`rollout`** (Task 2). The agent loop. Sample the model's tokens, hand each finished turn to
   the environment, append the tool's tokens, and record for every token who wrote it
   (`is_model`) and which turn it belongs to (`turn_id`). A test with a scripted "policy"
   checks the exact sequence `[<bos> ?] call 3 {<tool> higher} call 5 {<tool> lower} ans 4`
   from the chapter, and that you fed every token, tool output included, to the policy.
3. **`turn_advantages`** (Task 3). Credit for a group of G episodes of the same task:
   `"trajectory"` gives every turn `R_i - mean(R)`; `"turn"` gives turn k its reward-to-go
   minus the group baseline. Tests use the worked example from the chapter (Exercise 18.3).
4. **`masked_policy_loss`** (Task 4). REINFORCE on the tokens selected by the mask, averaged
   over those tokens. Tests check the value against the formula, and that positions predicting
   tool tokens and padding get *exactly* zero gradient.

The last test trains for 100 steps (about 10–20 s) and checks that the agent's success on
hidden-number tasks rises and that it calls the tool.

```bash
python -m pytest -q                   # all tests; they fail until you fill in the TODOs
python -m pytest -q -k task1          # just Task 1
RL_SOLUTIONS=1 python -m pytest -q    # run the tests against the reference solution
```

Dependencies: `numpy`, `torch`, `matplotlib`, `pytest`. CPU only.

## Run it

```bash
python run.py        # about 1.5-3 minutes; saves agentic_rl.png
```

It trains three variants with the same seed for 250 steps (32 tasks × 8 episodes per step),
each starting from a short supervised "cold start" that teaches the format but no strategy:

| Variant | What you should see (last 20 steps) |
|---|---|
| masked, trajectory-level | hidden number solved ≈ 0.9, about 2.5 tool calls: a near-binary search |
| masked, turn-level | about the same (the only per-turn rewards are small call costs) |
| **UN**masked | the agent abandons the tool; hidden number solved ≈ 0.1, the guessing rate is 1/8 |

The three panels: success on hidden-number tasks; tool calls per episode on revealed vs hidden
tasks (masked run); and tool calls on hidden tasks for all three variants. The script also
prints a few episodes of the trained agent. Read them: you should see it halve the interval.

On revealed tasks the agent learns to copy the number from the prompt, and it calls the tool
less there than on hidden tasks. But it drops those useless calls only slowly: about 1.4 calls
after 250 steps, about 1.0 after 600. A 0.05 cost changes the advantage by only ±0.05 when every
episode in the group succeeds, which is a weak signal next to the ±0.5 of hidden tasks.

## Things we learned building this lab (and what the chapter's exercises ask about)

- **A malformed turn must cost more than an honest failure.** With reward 0 for malformed
  turns, the agent learned to write garbage, because quitting at 0 beat a failed search at
  −0.15 (Exercise 18.6).
- **A tool that costs too much early is never learned.** With `call_cost=0.1`, all four seeds
  we tried stopped calling the tool within a few dozen steps. A random call rarely helps before
  the agent knows how to use the answer.
- **The sampler and the trainer must agree on what can be generated.** The `Policy` and the
  `Sampler` both forbid `<bos>` and `<pad>`. When only the sampler did, none of eight runs
  learned to use the tool well. Those two tokens' output weights, never corrected by sampling
  and pushed at full step size by Adam, drifted until at some positions the trainer put almost
  all its probability on them (Exercise 18.10 and Chapter 20).

## Going further (optional)

- Run `train(call_cost=0.1)` for a few seeds and confirm the collapse. Then add a cost
  *schedule*: 0 for the first 100 steps, rising linearly to 0.1 by step 200. Does the agent now
  learn the tool *and* stop calling it on revealed tasks?
- Divide each group's advantages by the group's reward standard deviation (GRPO's original
  normalization, Chapter 15). On revealed tasks, where every episode succeeds and rewards
  differ only by call costs, this inflates the cost differences. Does the agent stop calling the
  tool there faster? What does Chapter 16 say about this "difficulty bias"?
- Add a process reward of +0.02 for every valid tool call. What happens to the number of calls
  on revealed tasks? (This is the "paid for being busy" failure of the chapter's demo.)
- Measure the variance of the gradient as a function of `max_turns` (5, 10, 20) at a fixed
  policy, padding episodes with irrelevant turns. Compare with Exercise 18.4.