Two Loops, One Value Function: TD Learning Meets Reflexion

TL;DR. An agent has to cross a 32×32 grid full of randomly moving obstacles. Plain TD learning tunes the weights of a value function, but it can only tune the features a human chose. We wrapped it in an outer loop where an LLM reads the agent's recent experience, works out in plain language why the agent failed, and rewrites the value function's code. TD learning then fits the new weights from scratch. This is verbal reinforcement learning, in the spirit of Reflexion. In five iterations, the mean episode reward rose from −760 to about −116, against an optimum of −62. Every value function along the way was a few lines of readable Python.

The problem

The environment (env.py) is deliberately simple:

The shortest path is 62 steps, so the best achievable return is about −62. A random walk usually never reaches the goal and ends around −500 or worse.

The state is State(ai_position, goal_position, obstacle_positions). That is 66 integers, which is too many for a table but small enough that a few well-chosen features can capture what matters. The whole project rests on one question:

Which features?

Why the value function is the bottleneck

The agent doesn't learn a policy directly. It learns a parametric state-value function V(s; θ) and acts by one-step lookahead:

Q(s, a) = V(s')          where s' ~ P(· | s, a)    (one sampled successor)
a*      = argmax_a Q(s, a)   (ε-greedy during training)

The starting value function (value.py) is a reasonable-looking hand design:

V(s) = θ₁ · manhattan(ai, goal) + θ₂ · Σᵢ manhattan(ai, obstacleᵢ)

It has two parameters, is interpretable, and turns out to be badly wrong. The obstacle sum covers all 32 obstacles, so it is in the hundreds, and a single move changes it by up to ±32. The goal distance changes by only ±1. Worse, a sum of distances is a global quantity. It cannot tell "an obstacle is right next to me" from "obstacles are spread out", and that difference is exactly what predicts the −50 penalty.

No amount of gradient descent on θ fixes this. TD learning finds the best weights for the features you give it, and it can't invent new ones. With this V, the agent reaches a mean reward of −759.9 and mostly gets truncated at 500 steps while repeatedly hitting obstacles.

The fix needs a different function class, and that is where the outer loop comes in.

The big picture: two nested loops

flowchart TB
    subgraph OUTER["Outer loop: Reflexion / verbal RL (LLM)"]
        direction TB
        V0["value.py (version k)"] --> INNER
        subgraph INNER["Inner loop: TD learning (PyTorch)"]
            direction LR
            A["Act ε-greedily<br/>with V(s'; θ)"] --> B["Store transition<br/>in replay buffer"]
            B --> C["Sample minibatch,<br/>TD update on θ"]
            C --> A
        end
        INNER --> E["Evaluate: mean reward<br/>of last 10 episodes"]
        E --> S["Sample transitions from<br/>the last N episodes"]
        S --> R["LLM reflects in language:<br/>mechanics → failures →<br/>what V can't express → fix"]
        R --> VAL{"Validate code<br/>(imports, scalar,<br/>differentiable)"}
        VAL -- "error, retry ≤ 2" --> R
        VAL -- ok --> V1["value.py (version k+1)<br/>θ reset to 0"]
        V1 -.next iteration.-> V0
    end

The two loops search over different things.

The LLM never picks a number. It proposes features, such as "count of obstacles on my cell" or "obstacle one column to my right". TD learning then decides how much each one matters. The LLM is good at the symbolic part and bad at fitting weights; gradient descent is the opposite.

The inner loop: TD learning in detail

Acting

def q_values(value, env, state):
    # Q(s, a) = V(s') for ONE sampled successor per action
    return torch.stack([value(env.transition(state, a)[0]) for a in range(4)])

env.transition simulates a step without changing the environment. It samples obstacle moves and applies the agent's move. Actions are chosen ε-greedily, with ε decaying exponentially:

ε(frame) = 0.01 + (1.0 − 0.01) · exp(−frame / 10,000)

ε starts at 1.0 (fully random) and is about 0.14 at frame 20,000.

Learning

Every real step (s, a, r, s', done) goes into a replay buffer of 5,000 transitions. Once the buffer holds more than 1,000, each frame performs one semi-gradient Q-learning update on a minibatch of 32:

target = r                                  if done
       = r + γ · max_a' Q(s', a')           otherwise   (no gradient through the target)

loss   = mean over batch of (target − Q(s, a))²
θ ← Adam(θ, ∇_θ loss),  lr = 1e-2,  γ = 0.9

q_sa = value(env.transition(t.state, t.action)[0])  # differentiable
with torch.no_grad():
    target = t.reward + cfg.gamma * q_values(value, env, t.next_state).max()
loss = (target - q_sa).pow(2).mean()

Some details worth noting:

• Truncation is not termination. done is set only when the agent reaches the goal. An episode cut off at 500 steps still bootstraps, which is the correct treatment for time limits.
• The model is known but stochastic. Q(s,a) re-samples the obstacle moves each time it is evaluated, so the targets are noisy. Replay plus Adam smooths this out.
• Adam rather than SGD. The raw features are large (the original obstacle sum is in the hundreds), so plain SGD diverges unless lr ≲ 1e-6. Adam's per-parameter scaling makes the inner loop robust to whatever feature scales the LLM produces. This matters, because the outer loop shouldn't have to tune learning rates as well.
• θ starts at zero every iteration. Each new value.py gets a clean slate, so the result measures the function class, not leftover weights.

What the inner loop hands to the outer loop

train() returns a TrainResult with:

• the learned θ;
• per-episode rewards (the evaluation signal);
• the loss and θ history (saved to history.npz);
• the full transitions of the last N episodes (5 by default). This is the raw experience the LLM will reflect on.

The outer loop: Reflexion and verbal reinforcement learning

What Reflexion is

Reflexion (Shinn et al., 2023, "Reflexion: Language Agents with Verbal Reinforcement Learning") proposed that an LLM agent can improve without any weight updates. It does so by turning feedback from the environment into a written self-reflection that is stored in memory and conditions the next attempt. The paper describes three roles:

1. Actor. Produces behaviour (a trajectory).
2. Evaluator. Scores the behaviour (a reward, a unit test, a heuristic).
3. Self-Reflection model. Converts the score and trajectory into verbal feedback, such as "I failed because X, next time I should do Y".

The reflections act as a semantic gradient. The "parameters" being updated are text: a prompt, a plan, or here, source code. That is what verbal reinforcement learning means. The policy-improvement step happens in natural language rather than in weight space.

How this project maps onto Reflexion

One twist: in the original paper, reflection improves the actor's next attempt. Here, reflection improves the critic's function class, and an ordinary RL algorithm turns that critic into behaviour. The result is a hybrid: verbal RL on the structure, numerical RL on the parameters.

Step 1: choosing what the LLM sees

Five episodes can contain about 2,500 transitions, which is too many to send and mostly identical −1 steps. sample_transitions() reduces them to at most 100:

1. Find the most common reward (here −1). Everything else, such as −50, counts as rare.
2. Keep rare transitions, up to half the budget, evenly spaced if there are more.
3. Fill the remaining budget with evenly spaced common transitions.
4. Keep chronological order and tag each one with episode and t.

This is importance sampling for attention: failures are over-represented so the LLM has evidence to reason about, and enough ordinary steps remain to show loops and wall-bumping.

The LLM is told nothing about the task. There is no description of the grid, the obstacles, or the reward. It has to infer the mechanics from the transitions, which keeps the setup honest and general.

Step 2: structured reflection (the system prompt)

The prompt asks for reasoning before code, in four explicit steps:

1. Infer the mechanics. How do actions change the state? What does each reward mean? What moves on its own?
2. Find the bad transitions. Which ones have bad reward or no progress (loops, oscillation, repeated penalties)? What distinguishes them from the good ones?
3. Diagnose V. What can and can't the current V express about those features, and why does that lead to bad actions?
4. Make the smallest fix and sanity-check it on concrete transitions: would the new V rank the bad action below a better alternative, without harming the good transitions?

This turns a vague "make it better" into an explicit hypothesis → diagnosis → intervention → check cycle. Without it, in the earlier runs exp_5iter and exp_5iter_b, the LLM produced a short paragraph followed by code. Progress was slower and less stable (see "Results across all experiments" below).

Step 3: validation and retries

LLM-written code runs inside the training loop, so loading.py checks it before use:

• AST import allow-list: only torch, math, typing, __future__, and random_agent_rl;
• it must define class ValueFunction(ValueBase) with num_params;
• forward(state) must return a scalar tensor on sample states;
• backward() must give a finite gradient w.r.t. θ at θ = 0, so the inner loop can actually learn.

If a check fails, the error message is sent back to the LLM as the next user turn ("That value.py was rejected: …") for up to 2 retries. That is a small Reflexion loop of its own, where the evaluator is the validator. If the code never passes, the previous value.py is kept.

⚠️ The import allow-list is a guard rail, not a sandbox. Generated code still runs locally.

Step 4: hand-off

The validated file becomes proposal_{k+1}.py and the next iteration's value.py, with θ initialised to zeros. The newest proposal is always used; there is no rollback. That keeps the loop simple and lets you see regressions in the logs, but it is a design choice you may want to change (see "Lessons and practical tips" below).

A real run, iteration by iteration

This section follows the runs/exp_reasoning_5iter experiment (20,000 frames per iteration, seed 0, Claude, 100 transitions from the last 5 episodes).

Iteration 0: the hand-designed baseline → −759.9

V = θ₀·dist_goal + θ₁·Σ dist(ai, obstacle)
θ ≈ [-0.266, -0.013]

The LLM's reflection:

"Almost all of the −50 penalties happen when an obstacle sits at Manhattan distance 0–2 from the cell the agent steps into… The sum is a global quantity. One obstacle at distance 1 changes it by about 1 out of a total of several hundred… So V carries essentially no collision signal, and the −50 penalties cannot be represented."

It proposed replacing the sum with local proximity counts.

Iteration 1: local counts → −133.5 (5.7× better)

V = θ₀·dist_goal + θ₁·#obs@d0 + θ₂·#obs@d1 + θ₃·#obs@d2
θ ≈ [-0.22, -29.57, -2.30, -1.02]

Look at the learned weights. TD learning found by itself that sharing a cell with an obstacle is very costly (−29.6), being adjacent is mildly costly, and two cells away barely matters. The LLM supplied the right questions and TD supplied the answers. Episodes per iteration rose from 64 to 122, so the agent now reaches the goal regularly.

Iteration 2: a misleading sample → −146.9 (slight regression)

This time the sampled transitions happened to contain no −50 rewards. The LLM concluded, reasonably from what it saw, that obstacles "carry no visible reward effect". It dropped the distance-1 and distance-2 features and added a quadratic distance term to model discounting curvature.

That made the distance term non-monotone, so near the start, moving toward the goal sometimes looked no better than bumping into a wall.

Iteration 3: diagnosing its own mistake → −115.8 (best)

From new transitions, the LLM spotted the agent stuck at (0,0) for about 150 steps and traced the cause to its previous change:

"Its distance part is t0·d + t1·d², and the quadratic term lets it be non-monotone… moving toward the goal looks no better than staying put or bumping a wall, so the agent loops at (0,0). Make the distance term purely linear, so that reducing d always raises V."

It also noticed that obstacles move only horizontally and replaced the Manhattan-ball counts with same-row features:

V = θ₀·(d/31) + θ₁·#obs on my cell + θ₂·#obs in my row at col ±1
θ ≈ [-7.85, -25.36, -7.18]

This is the Reflexion pattern at its clearest: a failure, a verbal diagnosis that blames a specific earlier decision, and a targeted correction.

Here is that iteration-3 agent in action. The blue square is the agent, the red squares are the 32 obstacles, and the goal is the bottom-right corner. The agent makes its way through the drifting obstacles and reaches the goal in 82 steps (the optimum is 62).

Iteration 4: a plausible but wrong hypothesis → −131.9

The LLM claimed obstacles drift left "roughly 70–75%" of the time. It split the adjacency feature by side and added a right-at-distance-2 feature. The environment actually picks left or right uniformly, so the "asymmetry" was a pattern read into a small sample. The learned weights don't support it either: right-adjacent came out at −2.3 and left-adjacent at −5.5, the opposite of the hypothesis. Reward got slightly worse.

Lesson: verbal reflection can produce confident, well-argued, wrong stories. The inner loop and the evaluator are what keep it honest, because a feature that doesn't help gets a small weight and the score doesn't improve. Grounding matters.

Summary of the run

The reward is measured during training, so it still includes ε ≈ 0.14 exploration. The greedy policy does better than these numbers suggest.

Variant: critique cases, a regression suite, and long-term memory

The runs/exp_critique_5iter experiment (artifacts only; the default reflect.py uses the simpler loop above) adds two things that bring the loop closer to full Reflexion.

1. The LLM labels failures as well as rewriting code. Along with value.py, it returns a JSON list of critique cases:

{"id": 12, "bad_action": 2, "suggested_action": 1,
 "reason": "Moving to [29,30] collided with the obstacle at [29,29]; row+1 goes toward the goal and away from it"}

Each case (a state plus the action that should be preferred) is added to a growing regression suite (suite.json). After every training run, the greedy action of the new V is checked on every case (suite_eval.json):

2. Episodic memory across iterations. The prompt includes every earlier value.py together with its reward and suite pass rate, plus the failing suite cases with their Q-values. This is the long-term memory component of Reflexion. The LLM can see that "version 2 regressed to −549" and avoid repeating that idea.

The suite acts as a verbal unit test for the value function. It turns past reflections into checks that every future version has to answer to. Unit tests for code generation were one of the main evaluators in the Reflexion paper, and this variant applies the same idea to RL.

Results across all experiments

All runs share iteration 0 (same seed, same baseline, −759.9).

What the table shows:

• Every outer-loop variant beat the baseline by a wide margin. The function class was the bottleneck, not the optimiser.
• Structured reflection made progress faster and steadier. The 4-step prompt reached −133 on its first rewrite and never fell below −147. The unstructured runs swung by hundreds of points.
• Memory and tests help recover from regressions. The critique variant regressed at iteration 2, then recovered and kept improving once the history and suite showed it what had gone wrong.
• Progress is not monotone. Without rollback, a bad rewrite costs a whole iteration.

These numbers come from single runs with one seed. Treat them as illustrations, not benchmarks.

Lessons and practical tips

1. Split the work by strength. LLMs are good at structure: noticing that collisions are local, that obstacles move horizontally, that a quadratic term is non-monotone. Gradient descent is good at magnitudes. Don't ask the LLM for numbers, and don't expect SGD to invent features.
2. Make the reflection explicit. Asking for mechanics → failures → limitations of V → smallest fix → check, before the code, gave the most stable results here. It is chain-of-thought aimed at diagnosis.
3. Curate the evidence. What the LLM sees determines what it can conclude. Iteration 2 dropped the obstacle features because its sample happened to contain no collisions. Over-sample rare rewards, and consider guaranteeing at least a few examples of each reward type.
4. Expect confident hallucinations. The "70% left drift" story was detailed and wrong. Ground every hypothesis in measured outcomes: reward, suite pass rate, learned weights. A near-zero weight is the inner loop telling you a feature doesn't matter.
5. Validate generated code mechanically. Check that it parses, uses allowed imports, returns a scalar, and has finite gradients at θ = 0, and send the errors back to the LLM. Cheap checks catch most bad proposals before they waste 20,000 frames.
6. Reset θ when the features change. Old weights mean nothing for new features. Starting from zero also makes each iteration a fair test of the function class.
7. Use an optimiser that tolerates scale. Adam lets the outer loop propose features of any scale without retuning the learning rate.
8. Add memory and rollback for longer runs. Keeping version history, a regression suite, and "keep the best so far" (or "revert if reward drops by more than X") would turn the loop into a more reliable optimiser.
9. Interpretability comes for free. Every iteration's V is a short, documented Python file with a few weights, so you can read exactly what the agent has learned to care about.

Related ideas

• Reflexion (Shinn et al., 2023): verbal self-reflection and episodic memory as a substitute for weight updates. This is the main inspiration here.
• Eureka (Ma et al., 2023): LLMs iteratively write and refine reward functions from training feedback. This project does the same for the value function's features.
• FunSearch (Romera-Paredes et al., 2023): LLM-proposed programs, scored by an evaluator, refined in a loop.
• Classic feature engineering for linear function approximation (as in Sutton & Barto): the outer loop automates what a human RL practitioner does by looking at failure cases and adding a feature.

Glossary

• Value function V(s; θ): the expected discounted return from state s, here approximated by a hand- or LLM-designed function with a few weights θ.
• TD (temporal-difference) learning: updates V toward r + γ·V(next), bootstrapping from its own estimate instead of waiting for the episode to end.
• Semi-gradient: the gradient is taken through Q(s,a) only, not through the bootstrapped target.
• Experience replay: a buffer of past transitions that are sampled in random minibatches to break correlations between consecutive steps.
• ε-greedy: take a random action with probability ε, otherwise the best action according to Q.
• Function class / features: the form of V (which quantities it can see). The inner loop cannot change it; the outer loop can.
• Reflexion: a framework in which an LLM agent improves by writing verbal self-reflections on its failures and storing them in memory, without updating weights.
• Verbal reinforcement learning: reinforcement where the feedback and the update are expressed in natural language (and here, code) instead of gradients.
• Regression suite (critique variant): states labelled by the LLM with the action that should be preferred, which every new V is checked against.