Improving AI Search Agents With Co-evolving Feedback

# Illustrative sketch (not from the paper)
import torch

# 1. Load shared‑parameter model (e.g., Qwen2.5‑7B)
model = torch.nn.Module()  # placeholder for the actual LLM

# 2. Gather failure prefixes and create recovery demonstrations
def make_recovery_demo(failure_prefix):
    # keep erroneous prefix, append corrective feedback
    return failure_prefix + "[CORRECTIVE FEEDBACK]"

# 3. Online RL step with Comparative Feedback Estimate (CFE)
def online_rl_step(state):
    call_rollout = model(state)          # with feedback call
    skip_rollout = model(state)          # without feedback call
    # CFE = success_gap(call, skip) – abstracted as a scalar
    cfe = compute_success_gap(call_rollout, skip_rollout)
    # Feedback‑aware advantage shaping
    advantage = cfe * (call_rollout.log_prob - skip_rollout.log_prob)
    # Policy gradient update (simplified)
    loss = -advantage.mean()
    loss.backward()
    optimizer.step()
    optimizer.zero_grad()

# 4. Offline Rollout‑Derived Preference Optimization (RDPO) for the critic
def offline_rdpo(prefix_success, prefix_failure):
    # Preferential loss aligning critic scores with successful prefixes
    loss = preference_loss(critic(prefix_success), critic(prefix_failure))
    loss.backward()
    critic_optimizer.step()
    critic_optimizer.zero_grad()

# 5. Alternating training loop
optimizer = torch.optim.Adam(model.parameters(), lr=1e-5)
critic_optimizer = torch.optim.Adam(critic.parameters(), lr=1e-5)
for epoch in range(num_epochs):
    # online RL segment
    for state in online_batch:
        online_rl_step(state)
    # collect rollouts for RDPO
    success_prefixes, failure_prefixes = collect_prefixes()
    offline_rdpo(success_prefixes, failure_prefixes)

// Illustrative sketch (not from the paper)
const tf = require('@tensorflow/tfjs-node');

// 1. Load shared‑parameter model (placeholder for Qwen2.5‑7B)
let model = tf.sequential(); // actual LLM would be loaded here

// 2. Build recovery demonstrations from failure prefixes
function makeRecoveryDemo(failurePrefix) {
  // preserve erroneous prefix, append corrective feedback token
  return failurePrefix + '[CORRECTIVE FEEDBACK]';
}

// 3. Online RL step with Comparative Feedback Estimate (CFE)
async function onlineRlStep(state) {
  const callRollout = model.predict(state); // with feedback call
  const skipRollout = model.predict(state); // without feedback call
  const cfe = computeSuccessGap(callRollout, skipRollout); // scalar
  // Feedback‑aware advantage shaping
  const advantage = tf.mul(cfe, tf.sub(callRollout.logProb, skipRollout.logProb));
  const loss = tf.neg(tf.mean(advantage));
  const grads = tf.variableGrads(() => loss);
  optimizer.applyGradients(grads.grads);
  tf.dispose([loss, advantage, grads]);
}

// 4. Offline RDPO to update the critic
async function offlineRdpo(successPrefixes, failurePrefixes) {
  const posScore = critic.predict(successPrefixes);
  const negScore = critic.predict(failurePrefixes);
  const loss = preferenceLoss(posScore, negScore);
  const grads = tf.variableGrads(() => loss);
  criticOptimizer.applyGradients(grads.grads);
  tf.dispose([loss, grads]);
}

// 5. Alternating training loop
const optimizer = tf.train.adam(1e-5);
const criticOptimizer = tf.train.adam(1e-5);
for (let epoch = 0; epoch < NUM_EPOCHS; epoch++) {
  // online RL segment
  for (const state of onlineBatch) {
    await onlineRlStep(state);
  }
  // collect rollouts for RDPO
  const { successPrefixes, failurePrefixes } = collectPrefixes();
  await offlineRdpo(successPrefixes, failurePrefixes);
}