Reinforcement Learning for Code Agents

The formal framework behind self-learning agents: How reinforcement learning maps onto code generation – from Markov Decision Processes through Q-learning to the epsilon-greedy strategy.

1. From Feedback Loops to a Formal Framework

The previous article describes three pillars of self-learning agents: feedback loops, persistent memory, and self-critique. The feedback loop – agent generates, environment evaluates, agent adapts – follows a pattern formalized in AI research for decades: Reinforcement Learning (RL).

"The idea of RL is that an agent can learn from the environment by interacting with it and receiving rewards for performing actions." — Raschka et al., Machine Learning with PyTorch and Scikit-Learn, p. 1025

RL provides the mathematical structure for precisely describing feedback loops: Which states exist? Which actions are available? How is success measured? When is it worthwhile to try something new, and when to follow the proven strategy?

2. Markov Decision Process: The Model

A Markov Decision Process (MDP) formalizes the interaction between agent and environment through four components:

Mapping to Code Generation

Translating to a code agent yields a concrete MDP:

@dataclass(frozen=True)
class CodeState:
    task_type: str        # "bugfix", "feature", "refactor"
    complexity: int       # 1–5
    attempts: int         # Previous attempts
    test_pass_rate: float # 0.0–1.0

@dataclass(frozen=True)
class CodeAction:
    strategy: str         # "generate", "fix_errors", "refactor", "add_tests"
    context_level: int    # 1–3: minimal, normal, maximal

The Markov property states: the next state depends only on the current state and the chosen action – not on the entire history. For code agents, this is a simplifying assumption compensated by persistent memory (Article 01).

3. Reward Shaping: Differentiated Feedback

❌ Binary Feedback: Pass or Fail

A naive reward signal – 1.0 for passing tests, 0.0 for failure – provides insufficient information. The agent does not learn how close it was to a solution.

✅ Weighted Reward Signal

Reward shaping decomposes the evaluation into multiple metrics with different weights:

def compute_reward(
    test_pass_rate: float,
    quality_score: float,
    efficiency_score: float,
    weights: dict[str, float] | None = None,
) -> CodeReward:
    w = weights or {"test": 0.6, "quality": 0.25, "efficiency": 0.15}
    total = (
        w["test"] * test_pass_rate
        + w["quality"] * quality_score
        + w["efficiency"] * efficiency_score
    )
    return CodeReward(
        test_score=test_pass_rate,
        quality_score=quality_score,
        efficiency_score=efficiency_score,
        total=round(total, 4),
    )

The weighting reflects a clear priority: tests dominate (60%), code quality follows (25%), efficiency supplements (15%). These weights are configurable – a team can adjust the balance to its own standards.

Progress Bonus

In addition to the absolute score, a shaping term rewards progress relative to the previous state:

def shaped_reward(
    current_test_rate: float,
    previous_test_rate: float,
    base_reward: float,
    shaping_factor: float = 0.5,
) -> float:
    progress = current_test_rate - previous_test_rate
    return round(base_reward + shaping_factor * progress, 4)

An agent that improves the test pass rate from 0.3 to 0.8 receives a higher reward than one that remains constant at 0.8 – even though the absolute score is identical. This accelerates learning in early phases.

4. Q-Learning: Values for State-Action Pairs

The Q-Function

The central question in RL is: Which action is optimal in a given state? The Q-function Q(s, a) answers this question – it estimates the expected cumulative reward for action a in state s.

"The action value function determines the long-term rewards the agent will receive for taking action a. This function is usually expressed as Q(S, a), and hence is also called the Q-function." — Ashish Ranjan Jha, Mastering PyTorch, p. 393

The Bellman Equation

Q-values are not calculated directly but updated iteratively via the Bellman equation:

Q(s, a) ← Q(s, a) + α · [r + γ · max_a' Q(s', a') − Q(s, a)]

"The Bellman equation is one of the central elements of many RL algorithms. The Bellman equation simplifies the computation of the value function, such that rather than summing over multiple time steps, it uses a recursion." — Raschka et al., Machine Learning with PyTorch and Scikit-Learn, p. 1025

Implementation: Q-Table

For discrete state-action spaces, a Q-table stores values as a dictionary:

@dataclass
class QTable:
    values: dict[tuple[str, str], float] = field(default_factory=dict)
    default_value: float = 0.0

    def update(
        self, state_key, action_key, reward,
        next_state_key, actions, alpha=0.1, gamma=0.9,
    ) -> float:
        current_q = self.get(state_key, action_key)
        max_next_q = max(
            (self.get(next_state_key, a) for a in actions),
            default=self.default_value,
        )
        new_q = current_q + alpha * (reward + gamma * max_next_q - current_q)
        self.values[(state_key, action_key)] = round(new_q, 6)
        return new_q

    def best_action(self, state_key, actions) -> str:
        return max(actions, key=lambda a: self.get(state_key, a))

After sufficient episodes, Q-values converge: the agent "knows" which strategy yields the highest expected reward in each situation.

5. Exploration vs. Exploitation

The Dilemma

An agent that always chooses the action with the highest known Q-value (exploitation) may miss better strategies. An agent that only acts randomly (exploration) does not leverage its knowledge.

Epsilon-Greedy Strategy

The epsilon-greedy strategy resolves the dilemma with a simple mechanism:

• With probability ε: random action (exploration)
• With probability 1 − ε: best known action (exploitation)

def epsilon_greedy(
    q_table: QTable,
    state_key: str,
    actions: list[str],
    epsilon: float,
    rng: random.Random | None = None,
) -> tuple[str, bool]:
    r = rng.random() if rng else random.random()
    if r < epsilon:
        chosen = rng.choice(actions) if rng else random.choice(actions)
        return chosen, True
    return q_table.best_action(state_key, actions), False

"Under this mechanism, at each episode, an epsilon value is pre-decided, which is a scalar value between 0 and 1. [...] If the generated number is less than the pre-defined epsilon value, the agent takes a random action (exploration), otherwise it takes the action with the highest Q-value (exploitation)." — Ashish Ranjan Jha, Mastering PyTorch, p. 393

Epsilon Decay

Initially the agent explores heavily (high ε); with increasing experience, it relies more on learned knowledge (low ε). Exponential decay models this transition:

def decay_epsilon(
    epsilon: float,
    min_epsilon: float = 0.01,
    decay_rate: float = 0.005,
    episode: int = 0,
) -> float:
    decayed = min_epsilon + (epsilon - min_epsilon) * math.exp(-decay_rate * episode)
    return round(decayed, 6)

For code agents, this means: during initial tasks, various strategies (generate, refactor, fix errors) are tried. With growing experience, the optimal strategy for each task type crystallizes.

6. The Code-Agent-Environment Loop

All components together form the complete agent-environment loop:

    ┌─────────────┐
    │  CodeState   │ ◄───────────────────────────┐
    └──────┬──────┘                               │
           ▼                                      │
    ┌─────────────┐                               │
    │ ε-Greedy    │ ── Select action              │
    └──────┬──────┘                               │
           ▼                                      │
    ┌─────────────┐                               │
    │ Environment │ ── Execute code, run tests    │
    └──────┬──────┘                               │
           ▼                                      │
    ┌─────────────┐                               │
    │ Reward      │ ── Shaping + progress         │
    └──────┬──────┘                               │
           ▼                                      │
    ┌─────────────┐                               │
    │ Q-Update    │ ── Bellman equation            │
    └──────┬──────┘                               │
           ▼                                      │
    ┌─────────────┐                               │
    │ Next State  │ ─────────────────────────────►┘
    └─────────────┘

The run_episode function orchestrates a complete cycle:

def run_episode(
    initial_state, actions, q_table, environment,
    state_key_fn, epsilon=0.1, alpha=0.1, gamma=0.9,
    max_steps=10, rng=None,
) -> tuple[list[Transition], float]:
    transitions = []
    cumulative_reward = 0.0
    state = initial_state

    for step in range(max_steps):
        state_key = state_key_fn(state)
        action_key, _ = epsilon_greedy(q_table, state_key, actions, epsilon, rng)
        env_result = environment(state, action_key)

        reward_signal = compute_reward(
            env_result.test_pass_rate,
            env_result.quality_score,
            env_result.efficiency_score,
        )
        reward = shaped_reward(
            env_result.test_pass_rate,
            state.test_pass_rate,
            reward_signal.total,
        )

        next_state = CodeState(
            task_type=state.task_type,
            complexity=state.complexity,
            attempts=state.attempts + 1,
            test_pass_rate=env_result.test_pass_rate,
        )

        q_table.update(state_key, action_key, reward,
                       state_key_fn(next_state), actions, alpha, gamma)
        transitions.append(Transition(state, CodeAction(action_key, 1), reward, next_state))
        cumulative_reward += reward

        if env_result.is_terminal:
            break
        state = next_state

    return transitions, round(cumulative_reward, 4)

Each episode yields a list of transitions – a complete record of which state led to which action and what feedback the environment provided. These transitions form the foundation for experience replay (Article 04).

7. From RL to RLHF: The Bridge to LLMs

Reinforcement Learning from Human Feedback (RLHF) transfers RL concepts directly to the optimization of language models:

"RLHF combines reinforcement learning with human feedback to fine-tune language models. It aims to align the model's outputs with human preferences, improving the quality and safety of generated text." — Vaid & Lund, LLM Design Patterns, p. 452

The crucial difference: in classical RLHF, the model weights are adjusted. For code agents, the LLM remains unchanged – the "policy" is optimized through prompt adaptation, strategy selection, and context enrichment.

8. Conclusion

Reinforcement learning provides the formal scaffolding for self-learning code agents:

• MDP models the interaction as state → action → reward → next state
• Reward shaping transforms binary pass/fail into a differentiated signal
• Q-learning stores and updates the value of each strategy per situation
• Epsilon-greedy balances trying new approaches with exploiting proven strategies

The LLM itself is not trained. The "learning capability" emerges from the RL architecture around the model: the Q-table as knowledge store, epsilon decay as a maturation process, reward shaping as a differentiated feedback signal.

The next article in the series covers Automated Tests as Reward Signals – how test frameworks form the "environment" for code agents and how reward shaping is designed in practice.

Sources & Resources

• Raschka, S. et al. (2022): Machine Learning with PyTorch and Scikit-Learn. Packt. pp. 1025 ff. – Bellman equation, Q-learning fundamentals, policy concepts.
• Jha, A. R. (2024): Mastering PyTorch. 2nd Edition. Packt. pp. 393 ff. – Q-function, epsilon-greedy, Q-learning implementation.
• Elgendy, M. (2024): Modern Computer Vision with PyTorch. 2nd Edition. Packt. pp. 658 ff. – RL fundamentals, exploration vs. exploitation.
• Vaid, S. & Lund, B. (2025): LLM Design Patterns. Packt. pp. 452 ff. – RLHF, reward models, policy optimization for language models.
• Christiano, P. et al. (2017): Deep Reinforcement Learning from Human Preferences. NeurIPS. – Foundation of RLHF.
• Sutton, R. S. & Barto, A. G. (2018): Reinforcement Learning: An Introduction. 2nd Edition. MIT Press. – Standard reference for RL theory.

→ Full source code: src/