Self-Correcting Agents, Planning Agents and Custom LLM Training: Reflection Loops, Task Decomposition and LoRA

This is the last session of the course, and it covers the three ideas that turn an agent from a demo into something you can rely on. First, self-correction: checking its own work before answering. Second, planning: deciding the steps before taking them. Third, custom LLM training: what to do when the model's own behaviour has to change, which is the one place where you actually modify the weights.

Everything runs for real. The self-correction loop uses actual unit tests, and the fine-tuning trains a real model with a hand-built LoRA. Builds on agents (Session 28, 29) and on LoRA from Session 21.

Self-correcting agents

A self-correcting agent reviews its own work before it hands it back. The core is a reflection loop:

1. Generate an output.
2. Critique it against the goal.
3. Revise it.
4. Repeat until it passes or a retry limit is reached.

The most important design question is where the critique comes from. A model critiquing its own answer can simply approve a wrong one. The loop is far more powerful when the feedback is real and external: a failed unit test, a validator, a schema check, a compiler error. Real feedback is something the model cannot argue with. Here is a case where that matters. The task is deliberately ordinary-looking: write round_marks(x), which rounds exam marks "the way teachers round on paper". Our tests encode what teachers mean (2.5 becomes 3, 0.5 becomes 1), something the model has to discover.

import re, subprocess, sys, tempfile, textwrap

TASK = ("Write a Python function round_marks(x) that rounds a float of exam marks to a whole number, "
        "the way teachers round marks on paper. Reply with only the code.")
TESTS = textwrap.dedent('''
    cases = {1.4: 1, 1.6: 2, 2.5: 3, 3.5: 4, 0.5: 1, 10.0: 10, 74.5: 75, 99.49: 99}
    failed = [(k, v, round_marks(k)) for k, v in cases.items() if round_marks(k) != v]
    print(len(cases) - len(failed), "of", len(cases), "tests pass")
    for k, want, got in failed[:3]: print(f"  round_marks({k}) gave {got}, expected {want}")
''')

def extract_code(reply):
    m = re.search(r"```(?:python)?\n(.*?)```", reply, re.S)
    return (m.group(1) if m else reply).strip()

def check(code):
    """REAL external feedback: run the code against tests in a separate process."""
    with tempfile.NamedTemporaryFile("w", suffix=".py", delete=False) as f:
        f.write(code + "\n" + TESTS)
    r = subprocess.run([sys.executable, f.name], capture_output=True, text=True, timeout=10)
    output = (r.stdout + r.stderr).strip()
    return output.startswith("8 of 8"), output

code = extract_code(ask(TASK, max_tokens=350))
for round_ in range(1, 5):
    ok, feedback = check(code)
    print(f"round {round_}: {feedback.splitlines()[0] if feedback else 'no output'}")
    if ok: break
    for line in feedback.splitlines()[1:3]: print("   ", line)
    code = extract_code(ask(f"{TASK}\n\nYour previous code:\n```python\n{code}\n```\nIt failed these tests:\n{feedback}\nFix it. Reply with only the code.", max_tokens=350))
print("\nfinished after", round_, "round(s)", "- all tests pass" if ok else "- still failing, giving up at the limit")

round 1: 5 of 8 tests pass
      round_marks(2.5) gave 2, expected 3
      round_marks(0.5) gave 0, expected 1
round 2: 5 of 8 tests pass
      round_marks(2.5) gave 2, expected 3
      round_marks(0.5) gave 0, expected 1
round 3: 8 of 8 tests pass

finished after 3 round(s) - all tests pass

Watch the rounds. In round 1 the code passed 5 of 8 tests. The failures are exactly the half-way cases: round_marks(2.5) gave 2 and round_marks(0.5) gave 0. The model had used Python's built-in round, which rounds halves to the nearest even number ("banker's rounding"). That is correct by its own rules and wrong for a teacher. In round 2, given the failing tests, it still scored 5 of 8. In round 3 it found a fix and passed 8 of 8. The code never needed to understand why the tests failed. The test output carried the information, and the model used it.

That loop took three rounds on a 1.5-billion-parameter model, which is the realistic picture: self-correction often works, sometimes slowly, and sometimes not at all. That is why the loop has a limit: max_rounds caps the cost.

def self_correct(task, check, max_rounds=3):
    answer = ask(task)
    for _ in range(max_rounds):
        ok, feedback = check(answer)             # real feedback: run the tests, validate the JSON
        if ok:
            return answer
        answer = ask(f"{task}\n\nYour answer:\n{answer}\n\nIt failed:\n{feedback}\nFix it.")
    return answer                                # best effort after the limit

A good self-correction setup has four parts: error detection (the check), targeted repair (the feedback names what failed), state memory (the model sees its previous attempt) and budget control (the limit).

Planning agents: decompose first

A planning agent breaks a big goal into smaller, ordered sub-tasks before executing any of them. This is task decomposition. A goal like "build a weekly workout plan" becomes: fetch the user profile, estimate calorie needs, generate a split, then check it. Planning helps because a long task done in one go tends to lose track, while small steps can each be checked and, if needed, redone.

The common shape is planner and executor: one part writes the plan, another carries it out step by step, seeing the results so far. To re-plan after a failure, call the planner again with what has happened.

import re

goal = "Prepare a 3-day revision plan for FSc Physics chapter 'Current Electricity'."
raw = ask(f"Break this goal into 3 to 5 ordered steps, one per line, numbered. Steps only.\nGoal: {goal}", max_tokens=150)
plan = [re.sub(r"^\s*\d+[.)]\s*", "", line).strip() for line in raw.splitlines() if re.match(r"^\s*\d+[.)]", line)]
print("PLAN")
for i, step in enumerate(plan, 1): print(f"  {i}. {step}")

results = []
for i, step in enumerate(plan, 1):                          # the executor does one step at a time, seeing earlier results
    done = "\n".join(f"- {r}" for r in results)
    out = ask(f"Goal: {goal}\nAlready done:\n{done or '(nothing yet)'}\nNow do this step and answer in one short line: {step}", max_tokens=70)
    results.append(out)
    print(f"\nstep {i} result: {out}")

PLAN
  1. Identify key topics in 'Current Electricity' for FSc Physics.
  2. Plan daily study sessions for each topic, ensuring 3 days of revision.
  3. Allocate specific times for each topic, including review and practice problems.
  4. Include breaks and rest periods in the schedule to maintain focus.
  5. Review the entire chapter at the end of the 3 days to consolidate learning.

step 1 result: Identify key topics: Ohm's Law, Kirchhoff's Laws, DC circuits, series and parallel circuits, power in circuits, resistivity, and thermal effects of current.

step 2 result: Plan daily study sessions for Ohm's Law, Kirchhoff's Laws, DC circuits, series and parallel circuits, power in circuits, resistivity, and thermal effects of current, ensuring 3 days of revision.

step 3 result: Allocate 1 hour for Ohm's Law, 1 hour for Kirchhoff's Laws, 1 hour for DC circuits, 1 hour for series and parallel circuits, 1 hour for power in circuits, 1 hour for resistivity, and 1 hour for thermal effects of current.

step 4 result: Include 15-minute breaks every hour to maintain focus.

step 5 result: Review the entire chapter at the end of the 3 days to consolidate learning.

The planner turned "Prepare a 3-day revision plan for FSc Physics chapter 'Current Electricity'" into five ordered steps: identify the topics, plan daily sessions, allocate times, add breaks, and review. The executor then ran them one by one, each time seeing the earlier results. Step 1 produced a topic list (Ohm's law, Kirchhoff's laws and so on), and step 3 used that list to allocate an hour to each topic. The quality is that of a small model: sensible, a little generic, and step 2 mostly repeated the topics. A stronger model gives better plans, but the structure (plan first, then execute with context) is what makes long tasks manageable.

steps = ask("Break this goal into 3 to 6 ordered steps, one per line:\n" + goal)
plan = [s for s in steps.splitlines() if s.strip()]

results = []
for step in plan:
    results.append(ask(f"Goal: {goal}\nDone so far: {results}\nNow do: {step}"))

Custom LLM training: the last resort

Everything so far changed what the model sees (prompts, retrieved documents, tools, feedback). Training changes the model itself. It is the most expensive option, so go up the ladder one rung at a time:

The vocabulary of fine-tuning:

• Full training (pre-training): builds a model from scratch. Only for large labs.
• Fine-tuning: adapts an existing open model (LLaMA, Mistral, Qwen, Gemma) using a smaller curated dataset.
• SFT (supervised fine-tuning): training on input and output example pairs.
• PEFT: updating only a small part of the model.
• LoRA / QLoRA: PEFT methods with small add-on matrices. QLoRA adds 4-bit quantization so it fits on one GPU (Session 21).
• RLHF / DPO: aligning the model with human preferences (Session 13).

A real fine-tune, from scratch

Fine-tuning is best at teaching behaviour, not facts. Our task: a school grading rule that GPT-2 cannot know: 80 or above is A, 70 B, 60 C, 50 D, below 50 F. The prompt is Marks: 85. Grade: and the answer is a single letter. We measure on 30 marks the model never trained on, and compare before and after training LoRA matrices on two layers of every block of GPT-2 (small), with the LoRA written by hand in about 15 lines:

import random
import matplotlib.pyplot as plt
import torch
import torch.nn as nn
from transformers import AutoModelForCausalLM, AutoTokenizer

torch.manual_seed(0); random.seed(0)
tok = AutoTokenizer.from_pretrained("gpt2")
model = AutoModelForCausalLM.from_pretrained("gpt2")

# --- the task: a rule the base model has never been taught (our school's grading scale) ---
def grade(m): return "A" if m >= 80 else "B" if m >= 70 else "C" if m >= 60 else "D" if m >= 50 else "F"
make = lambda m: f"Marks: {m}. Grade:"
train = [(make(m), " " + grade(m)) for m in random.sample(range(0, 101), 70) for _ in range(3)]
test = [m for m in range(0, 101) if m not in {int(p.split()[1].strip('.')) for p, _ in train}][:30]
letters = [tok.encode(" " + g)[0] for g in "ABCDF"]

def accuracy():
    model.eval(); right = 0
    with torch.no_grad():
        for m in test:
            logits = model(**tok(make(m), return_tensors="pt")).logits[0, -1]
            right += "ABCDF"[int(logits[letters].argmax())] == grade(m)
    return right / len(test)

print(f"before training: {accuracy():.0%} of {len(test)} unseen marks graded correctly (guessing = 20%)")

# --- LoRA, written by hand for GPT-2's Conv1D layers (Session 21's idea, now on a real model) ---
class LoRAConv1D(nn.Module):
    def __init__(self, base, r=8, alpha=16):
        super().__init__()
        self.base = base
        for p in base.parameters(): p.requires_grad = False
        n_in, n_out = base.weight.shape                      # Conv1D stores (in, out)
        self.A = nn.Parameter(torch.randn(r, n_in) * 0.02)
        self.B = nn.Parameter(torch.zeros(n_out, r))
        self.scale = alpha / r
    def forward(self, x):
        return self.base(x) + self.scale * (x @ self.A.T) @ self.B.T

for p in model.parameters(): p.requires_grad = False
for block in model.transformer.h:
    block.attn.c_attn = LoRAConv1D(block.attn.c_attn)
    block.mlp.c_fc = LoRAConv1D(block.mlp.c_fc)
trainable = sum(p.numel() for p in model.parameters() if p.requires_grad)
total = sum(p.numel() for p in model.parameters())
print(f"trainable parameters: {trainable:,} of {total:,} ({trainable / total:.2%})")

opt = torch.optim.AdamW([p for p in model.parameters() if p.requires_grad], lr=3e-3)
losses = []
model.train()
for step in range(1, 241):
    batch = random.sample(train, 16)
    loss = 0
    for prompt, answer in batch:
        p_ids = tok.encode(prompt); a_id = tok.encode(answer)[0]
        logits = model(torch.tensor([p_ids])).logits[0, -1]
        loss = loss + nn.functional.cross_entropy(logits[None], torch.tensor([a_id]))
    loss = loss / len(batch)                                  # learn only the grade letter
    opt.zero_grad(); loss.backward(); opt.step()
    losses.append(loss.item())
    if step in (1, 60, 120, 240):
        print(f"step {step:3d}   loss {loss.item():.3f}")

print(f"after training : {accuracy():.0%} of {len(test)} unseen marks graded correctly")
model.eval()
for m in (91, 64, 38):
    out = model.generate(**tok(make(m), return_tensors="pt"), max_new_tokens=2, do_sample=False, pad_token_id=tok.eos_token_id)
    print(f"  {tok.decode(out[0])!r}   (correct: {grade(m)})")

fig, ax = plt.subplots(figsize=(6.2, 3.4))
ax.plot(losses, color="#DB2777", lw=1.5); ax.set_xlabel("training step"); ax.set_ylabel("loss")
ax.spines[["top", "right"]].set_visible(False)
fig.tight_layout()

before training: 7% of 30 unseen marks graded correctly (guessing = 20%)
trainable parameters: 663,552 of 125,103,360 (0.53%)
step   1   loss 3.348
step  60   loss 0.081
step 120   loss 0.000
step 240   loss 0.000
after training : 100% of 30 unseen marks graded correctly
  'Marks: 91. Grade: A A'   (correct: A)
  'Marks: 64. Grade: C C'   (correct: C)
  'Marks: 38. Grade: F F'   (correct: F)

Before training the model grades only 7 percent of the unseen marks correctly, below the 20 percent you get by guessing among five letters, since it has no idea of the rule. The LoRA adds 663,552 trainable parameters, 0.53 percent of the model's 125 million. The loss falls from 3.35 at step 1 to 0.08 at step 60 and to about zero by step 120. After training, it grades 100 percent of the 30 unseen marks correctly. It also generalises to marks it never saw (91 gives A, 64 gives C, 38 gives F), because it learned the rule, not a list.

Two honest caveats. This is a deliberately easy rule on a tiny dataset, and we judged only the five grade letters, so it shows the mechanism (cheap, small adapters teach new behaviour), not that real fine-tuning is always this easy. And it taught a rule, not facts: if you want the model to know your institute's fees, RAG is the better tool, which is why the ladder above puts it before fine-tuning. The generated text shows the letter twice (Grade: A A) because we trained only the first token after the prompt, a reminder that training data defines behaviour exactly, including quirks.

On real models, the libraries do this for you:

from peft import LoraConfig, get_peft_model
from trl import SFTConfig, SFTTrainer

base = AutoModelForCausalLM.from_pretrained("Qwen/Qwen2.5-1.5B-Instruct")
model = get_peft_model(base, LoraConfig(r=8, lora_alpha=16, target_modules=["q_proj", "v_proj"], task_type="CAUSAL_LM"))
trainer = SFTTrainer(model=model, train_dataset=my_pairs, args=SFTConfig(output_dir="out", num_train_epochs=3))
trainer.train()
model.save_pretrained("my-adapter")        # saves only the small adapter file

QLoRA loads the base in 4 bits first (BitsAndBytesConfig(load_in_4bit=True, bnb_4bit_quant_type="nf4")), which lets a 7-billion-parameter model be fine-tuned on one consumer GPU.

The whole path, in one page

Take a moment to look back. You started with a straight line fitted to a spreadsheet in Session 2. You learned how models are judged, tuned and regularised, how they actually learn by gradient descent and backpropagation, and why trees, forests, boosting and SVMs rule tabular data. You saw networks learn to see and to remember, met the Transformer that removed their limits, and watched it scale into language models. You gave those models knowledge (RAG), hands (tools and MCP), structure (frameworks and graphs), the ability to plan, remember and check themselves, and the option of training them.

The same five forces drove every step: straight lines were not enough (non-linearity), models either memorised or missed the pattern (bias versus variance), hand-made features hit a ceiling (learned features), sequential processing could not scale (parallelism), and a model alone could not know, act or check (tools, memory and feedback). If you remember the problem each idea solved, you will be able to place whatever new idea appears next.

Return to the roadmap and its mind-map story and read it again. It will read differently now.

Try it yourself

1. In self_correct_code.py, change the task to a rule of your own (for example, "round to the nearest 5 as shops do") and write the tests first. Does the loop find it?
2. In plan_execute.py, add a re-planning step: if a step's result is empty, call the planner again with the results so far.
3. In lora_finetune.py, change the grading rule (say, 85 and above is A) and retrain. How many steps does it need? Does 100 percent hold?
4. Fine-tune with only 5 distinct marks instead of 70. What happens to accuracy on unseen marks?

Where to go from here

Build something small and real: a question-answering bot over your own notes (RAG), then give it one tool, then add tracing and a test set. Read the original papers behind this course, starting with Attention Is All You Need. And keep the habit that ran through every session: measure, compare against a simple baseline, and be suspicious of a result that looks too good.

You now understand the machinery. The rest is practice.

Frequently asked questions