Skip to content

feat: BAVT budget-aware iteration controller (arXiv:2603.12634)#251

Closed
mlbrilliance wants to merge 3 commits into
huggingface:mainfrom
mlbrilliance:feat/bavt-budget-aware-agent
Closed

feat: BAVT budget-aware iteration controller (arXiv:2603.12634)#251
mlbrilliance wants to merge 3 commits into
huggingface:mainfrom
mlbrilliance:feat/bavt-budget-aware-agent

Conversation

@mlbrilliance

Copy link
Copy Markdown

Summary

Implements a training-free Budget-Aware Value Tree controller for the agent loop, adapted from:

"Spend Less, Reason Better: Budget-Aware Value Tree Search for LLM Agents"
Yushu Li, Wenlong Deng, Jiajin Li, Xiaoxiao Li — arXiv:2603.12634 (March 2026)

Why this paper?

ml-intern already has max_iterations (default 300), reasoning_effort, and a hash-based doom-loop detector — but no principled mechanism to adapt behaviour as the budget depletes. BAVT fills that gap without any training or LLM-call overhead.


Changes

agent/core/bavt.py (new)

Three classes, zero external deps:

Class Role
BudgetTracker Computes remaining/max ratio; suggests effort downgrade
ResidualProgressScorer Heuristic delta scorer over last 12 messages (no LLM calls)
BudgetConditionedController Combines both into a BudgetSignal per iteration

Budget thresholds:

budget_ratio Action
> 0.50 No action
0.25 – 0.50 Nudge if progress ≤ 0
0.10 – 0.25 Redirect message + effort downgrade
< 0.10 "Wrap up NOW" directive + max effort downgrade

5-iteration cooldown prevents back-to-back injections. max_iterations == -1 → ratio stays 1.0, all BAVT logic is a no-op.

agent/core/doom_loop.py

  • check_for_doom_loop() gains budget_ratio=1.0 param (backward-compatible)
  • Identical-call threshold tightens 3 → 2 at ratio < 0.25
  • detect_repeating_sequence() gains min_reps param; drops to 1 cycle at ratio < 0.10

agent/core/agent_loop.py

  • BudgetConditionedController instantiated once per run_agent() call
  • budget_ratio forwarded to check_for_doom_loop() each iteration
  • BudgetSignal.corrective_message injected as a user message (same pattern as existing doom-loop injection)
  • BudgetSignal.effort_hint forwarded to _resolve_llm_params() for dynamic effort downgrade

tests/test_bavt.py (new)

21 unit tests — all pass, no regressions in the full suite.

README.md

Architecture diagram updated + new "Budget-Aware Loop Control" section.


Test results

21 passed  ← new BAVT tests
353 passed, 3 skipped  ← existing suite (1 pre-existing failure in test_hub_artifacts unrelated to this PR)
ruff check: clean

Backward compatibility

  • No config changes required
  • --max-iterations -1 (unlimited): BAVT ratio stays 1.0, zero effect
  • check_for_doom_loop signature: new param has a default, all callers unaffected

Implements a training-free budget management system for the agent loop
inspired by 'Spend Less, Reason Better: Budget-Aware Value Tree Search
for LLM Agents' (Li et al., arXiv:2603.12634, March 2026).

## What this adds

### agent/core/bavt.py (new)
Three classes implementing the core BAVT logic:

- BudgetTracker: tracks remaining_iters/max_iters as a continuous ratio
  (1.0 = fresh, 0.0 = exhausted). When max_iterations == -1 (unlimited)
  the ratio stays at 1.0 so all BAVT logic is a no-op.

- ResidualProgressScorer: scores *relative* progress delta over the last
  12 messages using lightweight heuristics (no LLM calls, zero latency).
  Tracks unique content fingerprints + tool result quality to avoid the
  LLM self-evaluation overconfidence flaw the paper identifies.

- BudgetConditionedController: combines tracker + scorer into
  BudgetSignal output (corrective_message, effort_hint, prune_warning).
  A 5-iteration cooldown prevents back-to-back injections.

Budget thresholds:
  ratio > 0.50  → no action
  0.25-0.50     → nudge if progress stalled
  0.10-0.25     → redirect message + effort downgrade hint
  < 0.10        → 'wrap up NOW' directive + max effort downgrade

### agent/core/doom_loop.py
- check_for_doom_loop() gains a budget_ratio=1.0 parameter
- Identical-call threshold tightens 3→2 when budget_ratio < 0.25
- Sequence-repetition min_reps drops 2→1 when budget_ratio < 0.10
- detect_repeating_sequence() gains a min_reps parameter

### agent/core/agent_loop.py
- BudgetConditionedController instantiated once per run_agent() call
- budget_ratio computed from iteration / max_iterations each loop turn
- budget_ratio forwarded to check_for_doom_loop()
- BudgetSignal.corrective_message injected as a user message (same
  pattern as the existing doom-loop injection)
- BudgetSignal.effort_hint forwarded to _resolve_llm_params() so
  reasoning_effort is dynamically downgraded when budget is tight

### tests/test_bavt.py (new)
21 unit tests covering all three classes and key edge cases.

### README.md
Architecture diagram updated + new 'Budget-Aware Loop Control' section
with paper attribution.

Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>
Copilot AI review requested due to automatic review settings May 10, 2026 02:24

Copilot AI left a comment

Copy link
Copy Markdown

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Pull request overview

Adds a Budget-Aware Value Tree (BAVT)-style controller to make the agent loop adapt behavior (nudges, wrap-up directives, effort downgrades) as the iteration budget depletes, and tightens doom-loop detection under low budget.

Changes:

  • Introduces BudgetTracker, ResidualProgressScorer, and BudgetConditionedController to compute per-iteration budget/progress signals.
  • Makes doom-loop detection budget-aware by tightening thresholds based on budget_ratio.
  • Wires BAVT signals into run_agent() and documents the new control flow; adds unit tests for BAVT.

Reviewed changes

Copilot reviewed 5 out of 5 changed files in this pull request and generated 8 comments.

Show a summary per file
File Description
agent/core/bavt.py New BAVT controller + heuristics producing BudgetSignal each iteration.
agent/core/doom_loop.py Adds budget_ratio to doom-loop checks; introduces min_reps for sequence detection.
agent/core/agent_loop.py Instantiates BAVT controller, injects budget corrective messages, and forwards effort hints.
tests/test_bavt.py New unit tests for BAVT components.
README.md Updates loop diagram + adds “Budget-Aware Loop Control (BAVT)” section.

💡 Add Copilot custom instructions for smarter, more guided reviews. Learn how to get started.

Comment thread agent/core/doom_loop.py
Comment on lines 135 to 154
@@ -143,24 +150,37 @@ def detect_repeating_sequence(
else:
break

if reps >= 2:
if reps >= min_reps:
return pattern
Comment thread agent/core/agent_loop.py
Comment on lines +1180 to +1185
# Doom-loop detection: break out of repeated tool call patterns.
# Pass budget_ratio so thresholds tighten as iterations deplete.
budget_ratio = budget_controller._tracker.ratio(iteration)
doom_prompt = check_for_doom_loop(
session.context_manager.items, budget_ratio=budget_ratio
)
Comment thread agent/core/agent_loop.py
Comment on lines 1246 to 1251
llm_params = _resolve_llm_params(
session.config.model_name,
session.hf_token,
reasoning_effort=session.effective_effort_for(
session.config.model_name
),
reasoning_effort=bavt_effort_hint
or session.effective_effort_for(session.config.model_name),
)
Comment thread agent/core/bavt.py
Comment on lines +54 to +55
# Effort level ordering for downgrade logic (worst→best).
_EFFORT_ORDER = ["low", "minimal", "medium", "high", "xhigh", "max"]
Comment thread agent/core/bavt.py
Comment on lines +16 to +26
identifies as a core flaw of naive self-evaluation. This implementation is
entirely heuristic (no LLM calls) so it adds zero latency.

3. **Budget-conditioned signals** — thresholds on (budget_ratio, progress_score)
produce at most one injected message per check plus an optional effort hint
that the caller can forward to the next ``_resolve_llm_params`` call.

HuggingFace Pro note: The scorer uses only local heuristics by default. Callers
that want semantic scoring can optionally supply ``hf_token`` for a lightweight
HF Inference API call; this path is gated behind ``use_hf_scorer=True`` and
is never exercised in the default agent loop so it never adds latency.
Comment thread README.md
|---|---|
| ratio < 50 % and progress stalling | Inject a gentle nudge message |
| ratio < 25 % and progress negative | Inject redirect + downgrade `reasoning_effort` one level |
| ratio < 10 % | Inject a "wrap up now" directive + max effort downgrade |
Comment thread tests/test_bavt.py
Comment on lines +1 to +14
"""Unit tests for the BAVT (Budget-Aware Value Tree) module.

Covers BudgetTracker, ResidualProgressScorer, and BudgetConditionedController.
"""

from __future__ import annotations

from litellm import Message

from agent.core.bavt import (
BudgetConditionedController,
BudgetTracker,
ResidualProgressScorer,
)
Comment thread README.md
Comment on lines 264 to +274
╔═══════════════════════════════════════════╗
║ Iteration Loop (max 300) ║
║ ║
║ BAVT budget check (ratio + progress) ║
║ ↓ ║
║ Doom loop check (budget-aware) ║
║ ↓ ║
║ Get messages + tool specs ║
║ ↓ ║
║ litellm.acompletion() ║
║ (effort may be downgraded by BAVT) ║
copilot and others added 2 commits May 10, 2026 02:33
- tests/unit/test_bavt_integration.py: 17 behavioural tests covering all
  BAVT thresholds (nudge/downgrade/wrap-up), cooldown mechanics, doom-loop
  interaction, effort-hint flow-through, and residual scorer realism.
- tests/integration/test_live_bavt_hf.py: 4 live tests gated by
  ML_INTERN_LIVE_LLM_TESTS=1, proving end-to-end with real HF Inference API:
    1. effort_hint → correct HF router params (reasoning_effort header)
    2. real API call with BAVT-downgraded effort succeeds
    3. ResidualProgressScorer correctly ranks success vs error model output
    4. full stalling-agent scenario: BAVT fires → corrective message sent
       to real HF model → model responds with a coherent wrap-up pivot

All tests pass. Live tests verified against Qwen/Qwen2.5-72B-Instruct
via the HF Inference Router (https://router.huggingface.co/v1).

Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>
…3B proof

- Replace ml-intern streaming call with direct litellm.acompletion (no-tools)
  to avoid Qwen3/vllm rejection of empty tools=[] array
- Add _LLMResult dataclass + _live_call() helper (non-streaming, model-agnostic)
- Remove _mock_session (no longer needed), keep _IS_THINKING_MODEL + max_tokens
- All 4 live tests PASS on huggingface/Qwen/Qwen3.6-35B-A3B:
    test 1: effort_hint → valid HF router params (model, api_base, extra_body)
    test 2: real API call with BAVT effort='medium' → '2 + 2 equals 4.'
    test 3: ResidualProgressScorer success=1.00, error=-0.20 on real output
    test 4: end-to-end stalling → BAVT CRITICAL → LLM delivers Python scraper
- Add BAVT_PROOF_Qwen3.6-35B-A3B.txt evidence file (21s run, all green)

Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>
@mlbrilliance
mlbrilliance deleted the feat/bavt-budget-aware-agent branch May 10, 2026 03:02
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

2 participants