Agent Lightning 완벽 가이드

 

GitHub - microsoft/agent-lightning: The absolute trainer to light up AI agents.

 

---

Agent Lightning이란?

Agent Lightning은 Microsoft Research에서 개발한 오픈소스 프레임워크로, **강화학습(RL)**을 통해 AI 에이전트를 학습시킬 수 있는 혁신적인 도구입니다.

🎯 주요 특징

• 제로 코드 변경: 기존 에이전트 코드를 거의 수정하지 않고도 학습 가능
• 프레임워크 독립적: LangChain, AutoGen, OpenAI Agents SDK, CrewAI 등 모든 프레임워크 지원
• 선택적 최적화: 멀티 에이전트 시스템에서 특정 에이전트만 선택적으로 최적화 가능
• 다양한 알고리즘: 강화학습, 자동 프롬프트 최적화 등 다양한 학습 방법 지원

💡 왜 Agent Lightning인가?

기존의 AI 에이전트는 "정적"입니다. 한 번 만들어지면 그대로 사용됩니다. 하지만 Agent Lightning을 사용하면:

1. 실제 업무를 수행하며 학습: 에이전트가 작업을 수행하면서 동시에 개선됩니다
2. 복잡한 시나리오 처리: 다중 턴 대화, 도구 사용, 멀티 에이전트 협업 등 복잡한 상황에서도 학습 가능
3. 지속적인 개선: 더 많은 데이터로 계속해서 성능 향상

---

핵심 개념

1. 서버-클라이언트 아키텍처

Agent Lightning은 두 가지 핵심 컴포넌트로 구성됩니다:

┌─────────────────────────┐
│   Lightning Server      │  ← 학습 관리, 모델 제공
│  (학습 서버)             │
└─────────────┬───────────┘
              │
              │ 작업 전송 / 결과 수신
              │
┌─────────────▼───────────┐
│   Lightning Client      │  ← 에이전트 실행
│  (에이전트 실행 환경)     │
└─────────────────────────┘

Lightning Server (선생님 역할)

• 학습 데이터 관리
• 작업 샘플 준비
• LLM 엔드포인트 제공
• 모델 최적화 수행

Lightning Client (학생 역할)

• 서버로부터 작업 수신
• 에이전트 실행
• 실행 결과(궤적, trace) 서버로 전송

2. 마르코프 결정 과정 (MDP)

Agent Lightning은 에이전트 실행을 MDP로 모델링합니다:

• 상태(State): 현재 LLM에 입력되는 정보
• 행동(Action): LLM 호출
• 보상(Reward): 작업 성공/실패를 나타내는 신호

3. LightningRL

다중 단계 에이전트 실행을 단일 턴 RL로 변환하는 계층적 강화학습 알고리즘입니다.

---

설치 가이드

필수 요구사항

• Python 3.10 이상
• CUDA 지원 GPU (RL 학습 시)
• 가상환경 사용 강력 권장

1단계: 가상환경 생성

# conda 사용
conda create -n agent-lightning python=3.10
conda activate agent-lightning

# 또는 venv 사용
python -m venv agent-lightning-env
source agent-lightning-env/bin/activate  # Linux/Mac
# agent-lightning-env\Scripts\activate  # Windows

2단계: 필수 패키지 설치 (RL 학습용)

# PyTorch 설치
pip install torch==2.7.0 torchvision==0.22.0 torchaudio==2.7.0 \
    --index-url https://download.pytorch.org/whl/cu128

# FlashAttention 설치
pip install flash-attn --no-build-isolation

# vLLM 설치
pip install vllm==0.9.2

# VERL 설치
pip install verl==0.5.0

3단계: Agent Lightning 설치

pip install agentlightning

4단계: 에이전트 프레임워크 설치 (선택사항)

# AutoGen (권장: 먼저 설치)
pip install "autogen-agentchat" "autogen-ext[openai]"

# LangChain
pip install langgraph "langchain[openai]" langchain-community langchain-text-splitters

# OpenAI Agents
pip install openai-agents

# LiteLLM
pip install "litellm[proxy]"

# SQL 관련 (Spider 예제용)
pip install sqlparse nltk

참고: 의존성 충돌이 발생해도 괜찮습니다. 위의 설치 순서를 따르면 대부분 문제없이 작동합니다.

---

아키텍처 이해하기

실행 흐름

1. 작업 풀링 & 에이전트 실행
   ├─ Lightning Server가 작업 풀에서 작업을 가져옴
   └─ 에이전트로 작업 전송

2. 비침입적 추적 수집 (Sidecar 디자인)
   ├─ 에이전트가 작업 수행 (도구 사용, 멀티턴 대화 등)
   └─ 실행 추적, 에러, 보상 신호 자동 수집

3. 궤적 조직화 & 학습 루프
   ├─ 추적을 표준 전환 튜플로 변환: (state_t, action_t, reward_t, state_t+1)
   └─ LLM 최적화를 위한 학습 데이터 구조로 집계

통합 데이터 인터페이스

모든 에이전트 실행은 다음 형식으로 표현됩니다:

{
    "component": "에이전트 또는 도구 이름",
    "input": "입력 데이터",
    "output": "출력 데이터",
    "reward": 0.85  # 일반적으로 마지막 호출에만 있음
}

이 통합 인터페이스 덕분에:

• LLM 컴포넌트 외의 상태 변화도 캡처 가능
• 멀티 에이전트 시스템에서 특정 에이전트만 선택적 최적화 가능
• 프롬프트 튜닝 같은 모델 파인튜닝 외의 최적화도 지원

---

시작하기: 첫 번째 에이전트 만들기

간단한 질문-응답 에이전트를 만들어 자동 프롬프트 최적화(APO)를 적용해보겠습니다.

Step 1: 에이전트 클래스 생성

from agentlightning.litagent import LitAgent
from openai import OpenAI
import random

class SimpleAgent(LitAgent):
    def __init__(self):
        super().__init__()
        self.client = OpenAI()
    
    def training_rollout(self, task, rollout_id, resources):
        """단일 학습 롤아웃 실행"""
        
        # 1. task에서 질문 추출
        question = task["question"]
        
        # 2. resources에서 테스트할 시스템 프롬프트 가져오기
        system_prompt = resources["system_prompt"]
        
        # 3. OpenAI API 호출
        response = self.client.chat.completions.create(
            model="gpt-3.5-turbo",
            messages=[
                {"role": "system", "content": system_prompt},
                {"role": "user", "content": question}
            ]
        )
        
        answer = response.choices[0].message.content
        
        # 4. 보상 계산 (실제로는 정답과 비교하거나 LLM-as-judge 사용)
        # 여기서는 데모를 위해 랜덤 값 사용
        reward = random.uniform(0, 1)
        
        return reward

Step 2: 서버 설정 (server.py)

from agentlightning.server import AgentLightningServer

# 테스트할 시스템 프롬프트들
prompts = [
    "You are a helpful assistant.",
    "You are a knowledgeable AI.",
    "You are an expert assistant who provides detailed answers.",
]

# 테스트 작업들
tasks = [
    {"question": "What is the capital of France?"},
    {"question": "Explain quantum computing in simple terms."},
    {"question": "How does photosynthesis work?"},
]

# 서버 초기화
server = AgentLightningServer(
    tasks=tasks,
    host="127.0.0.1",
    port=9997
)

# 각 프롬프트에 대해 테스트
prompt_rewards = []

for prompt in prompts:
    print(f"\n[Algo] Testing prompt: '{prompt}'")
    
    # 작업 리소스에 현재 프롬프트 설정
    server.set_resources({"system_prompt": prompt})
    
    # 모든 작업 실행
    rewards = []
    for task in tasks:
        result = server.execute_task(task)
        rewards.append(result["reward"])
    
    avg_reward = sum(rewards) / len(rewards)
    prompt_rewards.append((prompt, avg_reward))
    print(f"[Algo] Received average reward: {avg_reward:.3f}")

# 최고 성능 프롬프트 찾기
best_prompt, best_reward = max(prompt_rewards, key=lambda x: x[1])
print(f"\n[Algo] Best prompt: '{best_prompt}' (reward: {best_reward:.3f})")

Step 3: 클라이언트 실행 (client.py)

from agentlightning.trainer import Trainer

# 에이전트 인스턴스 생성
agent = SimpleAgent()

# 트레이너 설정 (2개의 병렬 워커 사용)
trainer = Trainer(n_workers=2)

# 학습 시작
trainer.fit(agent, "http://127.0.0.1:9997")

Step 4: 실행하기

터미널 1 (서버):

python server.py

터미널 2 (클라이언트):

python client.py

예상 출력

[INFO] Initializing Trainer...
[INFO] Starting AgentOps local server on port 52081...
[INFO] Starting worker process 0...
[INFO] Starting worker process 1...

[Algo] Testing prompt: 'You are a helpful assistant.'
[INFO] [Worker 0 | Rollout] Completed in 1.09s. Reward: 0.631
[INFO] [Worker 1 | Rollout] Completed in 1.15s. Reward: 0.587
[Algo] Received average reward: 0.609

[Algo] Testing prompt: 'You are a knowledgeable AI.'
[INFO] [Worker 0 | Rollout] Completed in 1.22s. Reward: 0.925
[INFO] [Worker 1 | Rollout] Completed in 1.18s. Reward: 0.891
[Algo] Received average reward: 0.908

[Algo] Best prompt: 'You are a knowledgeable AI.' (reward: 0.925)

---

실전 예제

Agent Lightning 저장소에는 3가지 주요 예제가 포함되어 있습니다:

1. calc_x: AutoGen + Calculator Tool

목적: 계산기 도구를 사용하는 수학 문제 해결 에이전트

사용 기술:

• 프레임워크: AutoGen
• 데이터셋: Calc-X (GSM8K, Ape210K 기반)
• 기본 모델: Llama-3.2-3B-Instruct
• 학습 방법: 강화학습 (RL)

핵심 포인트:

• 언제 계산기를 호출할지 학습
• 중간 계산 값을 올바르게 통합하는 방법 학습

실행 방법:

cd examples/calc_x
python train_calc_agent.py

2. spider: LangGraph + SQL 생성

목적: SQL 쿼리를 작성하고 반복적으로 수정하는 에이전트

사용 기술:

• 프레임워크: LangGraph
• 데이터셋: Spider (Text-to-SQL)
• 기본 모델: Llama-3.2-3B-Instruct
• 학습 방법: 강화학습 (RL)

워크플로우:

1. Write (SQL 작성)
   ↓
2. Check (SQL 검증)
   ↓
3. Rewrite (필요시 수정)
   ↓
4. 반복 (최대 max_turns까지)

선택적 최적화:

• Write와 Rewrite 단계만 선택적으로 최적화
• Check 단계는 고정 (규칙 기반)

실행 방법:

cd examples/spider
python train_sql_agent.py

LangGraph 에이전트 구조 예시:

from langgraph.graph import StateGraph

class SQLAgentState(TypedDict):
    question: str
    sql_query: str
    validation_result: str
    iteration: int

def write_node(state):
    """SQL 쿼리 생성"""
    # LLM 호출하여 SQL 생성
    return {"sql_query": generated_sql}

def check_node(state):
    """SQL 검증"""
    # 실제 DB에서 실행해보고 에러 확인
    return {"validation_result": result}

def rewrite_node(state):
    """SQL 수정"""
    # 에러 메시지 기반으로 SQL 재작성
    return {"sql_query": revised_sql}

# 그래프 구성
workflow = StateGraph(SQLAgentState)
workflow.add_node("write", write_node)
workflow.add_node("check", check_node)
workflow.add_node("rewrite", rewrite_node)
workflow.add_conditional_edges("check", should_continue)

3. apo: 자동 프롬프트 최적화

목적: 커스텀 최적화 알고리즘 구현 방법 시연

사용 기술:

• 알고리즘: APO (Automatic Prompt Optimization)
• 최적화 대상: 시스템 프롬프트

핵심 개념:

• 모델 파인튜닝이 아닌 프롬프트 튜닝
• 여러 프롬프트 변형 테스트
• 최고 성능 프롬프트 자동 선택

실행 방법:

cd examples/apo
python train_apo_agent.py

---

고급 기능

1. 멀티 에이전트 시스템 최적화

멀티 에이전트 시스템에서 특정 에이전트만 선택적으로 최적화할 수 있습니다.

from agentlightning.litagent import LitAgent

class MultiAgentSystem(LitAgent):
    def __init__(self):
        super().__init__()
        self.planner = PlannerAgent()  # 최적화 대상
        self.executor = ExecutorAgent()  # 고정
        self.verifier = VerifierAgent()  # 최적화 대상
    
    def training_rollout(self, task, rollout_id, resources):
        # 1. Planner 실행 (최적화됨)
        plan = self.planner.create_plan(task)
        self.emit_llm_call("planner", plan)  # 추적 기록
        
        # 2. Executor 실행 (고정)
        result = self.executor.execute(plan)
        # emit 하지 않음 = 최적화하지 않음
        
        # 3. Verifier 실행 (최적화됨)
        verification = self.verifier.verify(result)
        self.emit_llm_call("verifier", verification)  # 추적 기록
        
        return self.calculate_reward(verification)

2. 커스텀 보상 함수

실제 작업 성공/실패를 반영하는 커스텀 보상 함수를 만들 수 있습니다.

def calculate_reward(self, task, result):
    """복합적인 보상 계산"""
    reward = 0.0
    
    # 1. 정확도 보상 (0.0 ~ 0.6)
    if self.is_correct(task["answer"], result):
        reward += 0.6
    
    # 2. 효율성 보상 (0.0 ~ 0.2)
    if result["num_steps"] <= 5:
        reward += 0.2 * (1 - result["num_steps"] / 5)
    
    # 3. 도구 사용 적절성 보상 (0.0 ~ 0.2)
    if self.used_tools_appropriately(result):
        reward += 0.2
    
    return reward

3. AIR (Automatic Intermediate Rewarding)

장기 태스크에서 희소 보상 문제를 해결하기 위한 중간 보상 자동 생성:

def training_rollout(self, task, rollout_id, resources):
    """중간 보상이 포함된 롤아웃"""
    
    total_reward = 0.0
    
    # Step 1: 도구 호출
    tool_result = self.call_calculator("2+2")
    if tool_result.success:
        intermediate_reward = 0.1  # 도구 사용 성공 보상
        self.emit_reward(intermediate_reward)
        total_reward += intermediate_reward
    
    # Step 2: 중간 결과 검증
    if self.validate_intermediate(tool_result):
        intermediate_reward = 0.2  # 검증 성공 보상
        self.emit_reward(intermediate_reward)
        total_reward += intermediate_reward
    
    # Step 3: 최종 답변 생성
    final_answer = self.generate_answer(tool_result)
    
    # 최종 보상
    final_reward = self.calculate_final_reward(task, final_answer)
    total_reward += final_reward
    
    return total_reward

4. 에러 모니터링 및 처리

class RobustAgent(LitAgent):
    def training_rollout(self, task, rollout_id, resources):
        try:
            result = self.execute_task(task)
            return self.calculate_reward(result)
        
        except TimeoutError:
            # 타임아웃 처리
            self.emit_error("timeout", rollout_id)
            return 0.0
        
        except ExecutionError as e:
            # 실행 에러 처리
            self.emit_error("execution_failed", rollout_id, str(e))
            return -0.5  # 페널티
        
        except Exception as e:
            # 예상치 못한 에러
            self.emit_error("unknown", rollout_id, str(e))
            return -1.0  # 더 큰 페널티

5. 추적 시각화

from agentlightning.trainer import Trainer

trainer = Trainer(n_workers=2)

# 에이전트 학습
trainer.fit(agent, "http://127.0.0.1:9997")

# 마지막 추적 시각화 (실험적 기능)
trainer.tracer.last_trace().visualize("tree_graph")

---

트러블슈팅

문제 1: "File not found" 에러

원인: Ray 초기화 시점의 환경 변수 및 작업 디렉토리가 중요합니다.

해결 방법:

# 현재 작업 디렉토리에서 Ray 재시작
ray stop
cd /path/to/your/project
ray start --head

문제 2: 학습 서버가 멈춤

원인: 에이전트 측에서 샘플 실패 또는 타임아웃

해결 방법:

# 프롬프트 및 응답 길이 제한 설정
trainer = Trainer(
    n_workers=2,
    max_prompt_length=2048,
    max_response_length=1024,
    timeout_seconds=60
)

문제 3: VERL OOM (Out of Memory)

원인: vLLM이 메모리 부족

해결 방법:

1. 체크포인트 자주 저장

server = AgentLightningServer(
    checkpoint_interval=100,  # 100 스텝마다 저장
    checkpoint_dir="./checkpoints"
)

1. 배치 크기 줄이기

server = AgentLightningServer(
    batch_size=4,  # 기본값보다 작게
    gradient_accumulation_steps=4  # 그래디언트 축적으로 보완
)

1. 더 작은 모델 사용

# Llama-3.2-3B 대신 Llama-3.2-1B 사용

문제 4: AgentOps 충돌

원인: 이미 코드에서 AgentOps를 사용 중

해결 방법:

trainer = Trainer(
    n_workers=2,
    tracer="custom",  # 관리형 AgentOps 비활성화
    disable_agentops=True
)

문제 5: 의존성 충돌

원인: 여러 프레임워크 설치 시 버전 충돌

해결 방법:

# 설치 순서를 정확히 따르세요
# 1. PyTorch, FlashAttention, vLLM, VERL
# 2. agentlightning
# 3. 에이전트 프레임워크들 (마지막)

# AutoGen을 먼저 설치
pip install "autogen-agentchat" "autogen-ext[openai]"

# 그 다음 다른 프레임워크들
pip install langgraph "langchain[openai]"

문제 6: 서버와 클라이언트 실행 순서

해결 방법: 일반적으로 순서는 중요하지 않지만, 다음 권장사항을 따르세요:

# 권장: 서버 먼저 시작
Terminal 1: python server.py
# 2-3초 대기
Terminal 2: python client.py

# 또는 백그라운드 실행
python server.py &
python client.py

---

성능 최적화 팁

1. 병렬 워커 수 조정

# CPU 코어 수에 따라 조정
import os
n_cores = os.cpu_count()
optimal_workers = min(n_cores - 2, 8)  # 최대 8개까지

trainer = Trainer(n_workers=optimal_workers)

2. 배치 처리 최적화

server = AgentLightningServer(
    batch_size=8,  # GPU 메모리에 맞게 조정
    prefetch_factor=2,  # 데이터 로딩 병렬화
)

3. 체크포인트 전략

# 빈번한 체크포인트 vs 성능
server = AgentLightningServer(
    checkpoint_interval=500,  # 너무 자주 저장하면 느려짐
    checkpoint_async=True,  # 비동기 저장으로 성능 개선
)

4. 로깅 레벨 조정

import logging

# 개발 중: DEBUG
logging.basicConfig(level=logging.DEBUG)

# 프로덕션: WARNING
logging.basicConfig(level=logging.WARNING)

---

추가 리소스

공식 문서 및 자료

• GitHub 저장소: https://github.com/microsoft/agent-lightning
• 공식 문서: https://microsoft.github.io/agent-lightning/
• Research Paper: Agent Lightning: Train ANY AI Agents with Reinforcement Learning
• Microsoft Research 프로젝트 페이지: https://www.microsoft.com/en-us/research/project/agent-lightning/

관련 프로젝트

• DeepWerewolf: AgentScope + Agent Lightning으로 구현된 중국 늑대인간 게임
• AgentFlow: Flow-GRPO 알고리즘을 사용하는 모듈형 멀티 에이전트 프레임워크

커뮤니티

• Reddit 토론: r/LocalLLaMA
• GitHub Issues: 버그 리포트 및 기능 요청
• Discord: (커뮤니티 채널 링크)

---

라이선스 및 기여

라이선스

MIT 라이선스 - 자유롭게 사용, 수정, 배포 가능

기여 방법

1. 포크 및 클론

git clone https://github.com/microsoft/agent-lightning
cd agent-lightning
pip install -e .[dev]

1. Pre-commit 훅 설치

pre-commit install
pre-commit run --all-files

1. 풀 리퀘스트 제출

• CLA (Contributor License Agreement) 동의 필요
• 코드 스타일 가이드 준수
• 테스트 작성

---

결론

Agent Lightning은 AI 에이전트를 "정적인 도구"에서 "학습하고 진화하는 시스템"으로 변화시키는 강력한 프레임워크입니다.

핵심 요약

✅ 쉬운 시작: 기존 에이전트에 몇 줄의 코드만 추가 ✅ 프레임워크 독립적: 어떤 에이전트 프레임워크와도 호환 ✅ 실전 최적화: 복잡한 멀티턴, 멀티 에이전트 시나리오 지원 ✅ 유연한 알고리즘: RL, APO, SFT 등 다양한 학습 방법

다음 단계

1. 간단한 예제부터 시작: examples/apo로 프롬프트 최적화 실험
2. 도구 사용 학습: examples/calc_x로 에이전트가 도구를 효과적으로 사용하도록 학습
3. 복잡한 워크플로우: examples/spider로 멀티 스텝 SQL 에이전트 구축
4. 커스텀 에이전트: 실제 업무에 맞는 에이전트 개발 및 최적화

Agent Lightning과 함께 여러분의 AI 에이전트를 더욱 똑똑하게 만들어보세요! 🚀⚡