Agent Lightning Integration

Bridge Agent Lightning rollout rewards into Mubit. Rewards flow back to the recalled lessons that shaped each rollout via record_outcome.

Agent Lightning is an RL/optimization framework for agents: it runs rollouts and credits reward back to whatever shaped them. That makes it a natural fit for Mubit's attribution loop — a reward should credit the recalled lessons that influenced a rollout, not just the run as a whole. The mubit-agent-lightning adapter provides two surfaces:

• MubitHook — a subclass of agentlightning.Hook. Register it on a Trainer and the loop runs itself: recall at rollout start, credit at rollout end.
• MubitRewardEmitter — an explicit emitter for custom training loops or inside LitAgent.rollout, with reward, event, and step-outcome methods.

pip install "mubit-agent-lightning[agent-lightning]"

This pulls mubit-sdk>=0.9.0 and agentlightning>=0.3,<0.4.

Minimal usage

import agentlightning as agl
from mubit_agent_lightning import MubitHook
 
hook = MubitHook(
    endpoint="https://api.mubit.ai",
    api_key="mbt_...",
    session_id="my-rl-run",
    recall_query="",              # empty -> grounds recall in each rollout's input
    verified_in_production=True,  # flag lessons confirmed in live use for a boost
)
 
trainer = agl.Trainer(..., hooks=[hook])
trainer.fit(my_lit_agent, ...)

No wiring inside the agent is required — the hook rides the rollout lifecycle.

How the hook closes the loop

• on_rollout_start — recalls relevant Mubit lessons (using recall_query, or the rollout's own input text when the query is empty) and stashes the recalled entry_ids plus grounding citations keyed by rollout_id.
• on_rollout_end — pulls the rollout's final reward out of the collected spans (agl.find_final_reward) and forwards it to record_outcome, crediting the stashed entry_ids. Reward >= 0 records a success outcome, negative records a failure; the raw value is carried as the signal.

Both lifecycle methods fail open: a Mubit hiccup never breaks a rollout. On the next recall, the lessons that earned reward rank higher — see The learning loop.

Full example: Explicit emitter loop

For custom training loops, MubitRewardEmitter gives you each step of the loop explicitly: recall, act, credit.

import os
 
from mubit_agent_lightning import (
    MubitRewardEmitter,
    extract_citations,
    extract_entry_ids,
)
 
TASK = "What is the safe way to retry a failed tool call?"
 
emitter = MubitRewardEmitter(
    endpoint=os.environ.get("MUBIT_ENDPOINT", "https://api.mubit.ai"),
    api_key=os.environ["MUBIT_API_KEY"],
    session_id="agent-lightning-demo",
    agent_id="reward-rollout-demo",
)
 
# --- recall: pull memory to ground the rollout ---
recall = emitter.recall(TASK, limit=10)
entry_ids = extract_entry_ids(recall)    # IDs of every recalled lesson
citations = extract_citations(recall)    # which ones grounded the answer
 
# --- act: your agent solves the task using recall["evidence"] ---
evidence = recall.get("evidence", [])
reward = run_rollout(TASK, evidence)     # your rollout; returns a scalar reward
 
# --- intra-rollout (process) credit for a single step ---
emitter.emit_step_outcome(
    "retrieval",
    0.8,
    step_name="recall-and-ground",
    rationale="recall returned usable grounding for the task",
    directive_hint="prefer reference_id lessons over raw traces",
    entry_ids=entry_ids,
)
 
# --- close the loop: reward flows back to the recalled lessons ---
emitter.emit_reward_for_recall(
    reward,
    recall,
    verified_in_production=True,
    rationale="task solved using recalled lessons",
)
 
# --- durable event with server-side dedup on retry ---
emitter.emit_event(
    "rollout_summary",
    f"solved {TASK!r} with reward {reward}",
    idempotency_key="demo:rollout-001:event-1",
)

emit_reward_for_recall(reward, recall) is shorthand for emit_reward(reward, entry_ids=extract_entry_ids(recall)); pass cited_only=True to credit only the evidence the answer was actually grounded in.

Emitter surface

emitter.recall(query, limit=10, entry_types=None)      # -> raw recall dict
emitter.emit_reward(reward, entry_ids=[...],           # -> record_outcome
                    verified_in_production=True,
                    reference_id=None, rationale="")
emitter.emit_reward_for_recall(reward, recall_result,  # recall + credit in one call
                               cited_only=False)
emitter.emit_event(kind, content, intent="trace",      # -> remember
                   idempotency_key=None)
emitter.emit_step_outcome(step_id, signal,             # -> record_step_outcome
                          step_name="", entry_ids=[...],
                          directive_hint=None)

Re-exported helpers: extract_entry_ids(result, cited_only=...), extract_citations(result), normalize_metadata(evidence). MubitTracer is retained as a backward-compatibility alias of MubitRewardEmitter (adds the legacy trace_reward(content, reward) convenience). See step-level outcomes for how per-step signals are used at recall time.

Configuration

Gotchas

• Reward sign picks the outcome label. reward >= 0 records success, negative records failure; the exact float is preserved as signal. Shape your reward accordingly if you use zero as "neutral".
• reference_id defaults to the run-level "global" sentinel. You don't have to map every reward to a single lesson — entry_ids carries the per-entry attribution.
• mirror_to_agl=True double-reports by design. The emitter writes to Mubit and to Agent Lightning's own span store so the trainer's reward accounting still sees the value. Set it to False in loops where the framework already captured the reward.
• Everything fails open. Recall or record errors inside MubitHook are swallowed so a Mubit outage cannot break training.
• Use idempotency_key on retried emit_event writes. At-least-once delivery with the same token dedupes server-side instead of double-counting reinforcement.

Version compatibility

The Hook ABC (on_rollout_start / on_rollout_end), the LitAgent rollout API, and the module-level emit_reward / emit_message / find_final_reward helpers the adapter conforms to were verified against agentlightning 0.3.0.