DOCSEARCHbyAlgolia
Submit Docs

Agent Lightning

repository·main·Indexed 12 days ago

https://github.com/microsoft/agent-lightning

A framework-agnostic trainer for optimizing AI agents using Reinforcement Learning (RL) and Automatic Prompt Optimization. It supports agents built with LangChain, AutoGen, OpenAI, and other SDKs. Version 0.3.1 includes support for algorithms like GRPO and EMPO², as well as integrations for kernel-level safety via Agent-OS and environments such as ALFWorld and ScienceWorld.

Tokens
85.1K
Snippets
249
Records
373
Agent score
96%

What's inside Agent Lightning

  2. Overview of Agent Lightning

    main

    Agent Lightning is a training framework designed to optimize AI agents using various algorithms such as Reinforcement Learning (RL), Automatic Prompt Optimization (APO), and Supervised Fine-tuning (SFT).

    Key capabilities include:

    • Framework Agnostic: Works with LangChain, OpenAI Agent SDK, AutoGen, CrewAI, Microsoft Agent Framework, or even raw Python OpenAI implementations.
    • Zero Code Change (almost): Designed to turn existing agents into optimizable entities with minimal modifications.
    • Selective Optimization: Allows for the optimization of specific agents within a larger multi-agent system.
    • Algorithm Support: Provides built-in support for Reinforcement Learning and prompt optimization techniques.

  3. Explore the Agent-lightning Examples Catalog

    main

    The Agent-lightning examples catalog provides a variety of reference implementations ranging from minimal primitives to complex reinforcement learning workflows. You can use these examples to understand how to implement specific agent behaviors, integrate with external tools, or perform fine-tuning.

    Key categories of examples include:

    • Agent Optimization & RL: Includes APO room selector (prompt optimization), Calc-X VERL math (reinforcement learning with AutoGen), ChartQA vision-language RL (LangGraph + VERL/GRPO), and Spider SQL agent (text-to-SQL training).
    • Fine-tuning (SFT) Workflows: Includes Azure OpenAI SFT (rollout to JSONL to Azure fine-tune) and Unsloth SFT (4-bit LoRA fine-tuning using Unsloth).
    • Integration & Tooling: Includes Claude Code SWE-bench (instrumented driver for Anthropic workflows), Tinker integration (agl_tinker adapter for CrewAI/OpenAI), and RAG (MuSiQue) pipelines.
    • Core Primitives: The minimal examples provide bite-sized scripts for studying LightningStore, LLM proxying, and minimal vLLM hosting.

  4. Explore Agent-lightning examples and recipes

    main

    Agent-lightning provides a catalog of examples covering various agentic workflows, training methods, and integrations. You can use these examples to understand how to implement specific patterns or to see how different building blocks interact.

    Core Example Categories:

    • Optimization & Fine-tuning:
      • apo: Automatic Prompt Optimization (built-in, custom, and debugging workflows).
      • azure: Supervised fine-tuning using Azure OpenAI.
      • unsloth: Supervised fine-tuning using Unsloth with 4-bit quantization and LoRA.
    • Agent Training & Reasoning:
      • calc_x: Math reasoning agent training using VERL and AutoGen with an MCP calculator tool.
      • chartqa: Vision-language ChartQA agent using LangGraph, VERL, and multi-step self-refinement.
      • spider: Text-to-SQL reinforcement learning training on the Spider dataset using LangGraph.
      • tinker: Reinforcement learning using Tinker as the backend training service.
    • Specialized Pipelines:
      • rag: Retrieval-Augmented Generation pipeline targeting the MuSiQue dataset with Wikipedia retrieval.
      • claude_code: SWE-bench harness for Claude Code that records traces across Anthropic, vLLM, and OpenAI-compatible backends.
    • Learning Building Blocks:
      • minimal: Small programs demonstrating individual Agent-lightning components in isolation.

    Community-contributed examples and recipes can be found in the ../contrib directory.

  5. Use built-in algorithms in AgentLightning

    main

    AgentLightning provides a library of pre-built algorithms designed for compatibility across various agent scenarios. You can use these algorithms directly within your agent workflows.

    Currently available algorithms include:

    • APO (Automatic Prompt Optimization): Uses textual gradients and beam search to optimize prompts.
    • VERL: Implements Reinforcement Learning using the VERL framework.

    To customize these algorithms, refer to the Algorithm-side References.

  6. What is the LightningStore and how does it work?

    main

    The LightningStore is the central coordination point for Agent-lightning. It manages the lifecycle of work through several key abstractions:

    • Task Queue: Managed via enqueue_rollout (to add work) and dequeue_rollout (for workers to poll). Dequeuing a rollout automatically creates its first associated attempt.
    • Rollouts: The high-level unit of work. A rollout represents the overall task and has a lifecycle (e.g., queuing → preparing → running → succeeded). Algorithms and Runners typically interact with the rollout as a single entity.
    • Attempts: The "inside view" of a rollout. Each rollout can have multiple attempts (retries). Attempts track execution details like status, start_time, end_time, and last_heartbeat_time.
    • Spans: Structured trace events produced during an attempt, ordered by a monotonic sequence ID per (rollout_id, attempt_id).
    • Resources: Versioned, named bundles (like prompt templates) referenced by rollouts.
    • Workers: Metadata about runner instances, including heartbeat timestamps and current assignments.

    Mental Model: Think of Rollouts as the external view of a task and Attempts as the internal execution attempts of that task. A rollout's status is an aggregated view of its latest attempt's status plus control actions like cancellation or queueing.

  7. Configure Execution Strategies in agentlightning

    main

    Execution strategies define how training tasks are distributed and executed. You can choose between different strategies depending on your hardware and concurrency requirements:

    • agentlightning.ExecutionStrategy: The base class for all execution strategies.
    • agentlightning.ClientServerExecutionStrategy: Suitable for distributed setups where a client communicates with a central server to manage tasks.
    • agentlightning.SharedMemoryExecutionStrategy: Optimized for single-node training where processes share memory to reduce overhead.
    from agentlightning import ClientServerExecutionStrategy, SharedMemoryExecutionStrategy

# Example: Using Shared Memory for local high-performance training
strategy = SharedMemoryExecutionStrategy()

  8. How Spans work in the Store

    main

    Every traceable operation in a rollout is stored as a Span. Spans serve two purposes:

    1. Instrumentation: They capture fine-grained trace data.
    2. Heartbeats: They act as periodic heartbeats to demonstrate liveness. The first span marks activation, and subsequent spans refresh the attempt's last_heartbeat_time.

    Watchdog Behavior: If no span is received within the configured unresponsive_seconds (defined in RolloutConfig), the watchdog downgrades the attempt status to unresponsive until new activity (a new span) resumes.

    Ordering: Spans are indexed by (rollout_id, attempt_id, sequence_id). To ensure correct ordering in distributed systems despite clock skew, the store enforces a monotonically increasing sequence_id for every span created within an attempt.

    // Spans are indexed by:
// (rollout_id, attempt_id, sequence_id)

  9. Understand the Azure OpenAI Fine-tuning Workflow

    main

    The fine-tuning loop follows a four-stage automated process:

    1. Collect traces: The Trainer uses runners to gather agent rollouts against a base deployment in batches defined by finetune_every_n_rollouts.
    2. Filter and package data: Agent-lightning collects rewards and telemetry. It uses data_filter_ratio to filter traces, then serializes the remaining high-quality traces into Azure OpenAI JSONL format.
    3. Fine-tune: The AzureOpenAIFinetune.finetune method uploads the dataset to Azure, waits for the job to complete, and returns the new model identifier.
    4. Deploy and evaluate: A new versioned deployment (e.g., gpt-4.1-mini-ft_v01) is created. If the number of deployments exceeds max_deployments, older ones are pruned. Finally, validation rollouts are run to confirm the reward of the new model.

  10. Understand VERL reward propagation and customization

    main

    The VERL wrapper decomposes agent executions into prompt–response pairs via an Adapter and associates them with reward signals as Triplet objects.

    Reward Strategy: The final scalar reward (derived from the last triplet in a trajectory) is propagated to all preceding triplets in that trajectory. This ensures every triplet receives an identical reward signal, allowing them to be optimized as valid RLHF trajectories.

    Customization: Currently, fine-grained control over reward propagation or credit assignment is not exposed via the public API. If you require customized reward shaping or trajectory decomposition, you must clone and modify the VERL source implementation directly.

  12. How bundles and execution strategies work together

    main

    Agent-lightning scales by splitting training into two distinct bundles coordinated by an execution strategy:

    1. Algorithm Bundle: Wraps your Algorithm, Adapter, and LLM proxy. It is a single callable that can be aborted via a signal event. Its signature is: async def algorithm_bundle(store: LightningStore, event: ExecutionEvent) -> None

    2. Runner Bundle: Wraps the Runner, Tracer, hooks, and agent. Unlike the algorithm bundle, runner bundles are intended to be replicated across multiple workers. Its signature is: async def runner_bundle(store: LightningStore, worker_id: int, event: ExecutionEvent) -> None

    An execution strategy determines where these bundles are placed (threads, processes, or different machines), how many runner replicas to launch, and how to coordinate lifecycle events like shutdown.

    By default, the Trainer uses an InMemoryLightningStore. The execution strategy wraps this store in thread-safe or HTTP-safe facades (like LightningStoreThreaded or LightningStoreServer) depending on the placement of the bundles.