[논문리뷰](25.01) Agent-R: Training Language Model Agents toReflect via Iterative Self-Training

DeepLearing/NLP(Agent)

[논문리뷰](25.01) Agent-R: Training Language Model Agents toReflect via Iterative Self-Training

notdecidedyet 2025. 4. 19. 14:03

Monte Carlo Tree Search(MCTS)를 활용하여 다양한 궤적(좋은 궤적/나쁜 궤적)을 탐색하고 수집
액터 모델이 나쁜 궤적 내 첫 번째 오류를 식별한 후, 해당 지점에서 좋은 궤적과 연결하는 모델 가이드 수정 궤적(revision trajectory)을 구성
이렇게 구성된 수정 궤적을 사용하여 모델을 반복적으로 훈련시킴으로써, 에이전트가 오류를 실시간으로 감지하고 자체 반영하는 능력을 향상시

Loss:
- 윗줄에 있는 텀: 궤적 데이터 - SFT - 좋은 궤적은 다 학습하고, 수정 궤적은 수정 이후의 궤적만 학습함.
  -> 수정 궤적의 앞부분은 틀린 정보를 포함할 수 있기 때문에 학습하지 않는 것으로 보임.
- 밑줄에 있는 텀: ShareGPT데이터 셋

Abstract

배경:

LLM 에이전트들은 상호작용적 환경에서 중요한 역할을 하지만, 기존 연구들은 스스로 오류 회복하는 능력 부족함
이를 획득하기 위한 step-level critique 데이터는 수집하기 어렵고 비용이 많이 듦

논문 제안:

Agent-R: LLM이 Reflect할 수 있는 반복적 훈련 프레임워크 제안
- MCTS를 활용하여 틀린 궤적에서 올바른 궤적으로 훈련 데이터 구성
- 핵심 내용: 롤아웃의 끝까지 가지 않고 중간에 수정하는 능력을 갖춤.
  - 이를 위해 model guide critique 메너니즘 도입: 액터 모델이 틀린 궤적에서 오류를 식별하고 올바른 경로와 연결

Intropduction

기존 연구 및 한계:

- LLM 에이전트들은 자율적 의사 결정, 오류 수정, 작업 최적화 능력이 필요한 환경에서 활용됨
- 기존 방법들은 강력한 전문가(human/ai)의 행동을 따라하나, 이는 실제 환경에서 실패하기 쉬움

주요 문제: 오류 감지와 동적 궤적 수정 능력 부족

기존 접근법의 한계:

기존 연구의 한계를 극복하기 위해, 명시적 오류 신호나 보상 함수에 의존하는 방법들을 사용함.
-> 허나 코드 수정, 수학 추론과 같은 single 턴 시나리오에만 초점
-> 상호작용 환경에서 제안한 연구는 없음
-> 또한 상호작용 환경은 긴 추론이 필요한데, 높은 품질의 보상 함수를 설계하는것은 어려움.

제안하는 해결책 - Agent-R:

틀림을 인지하는 즉시 개선하는 동적 자체 훈련 프레임워크 제안
- MCTS를 활용하여 잘못된 궤적에서 올바른 궤적으로 복구하는 프로세스
- 가장 적합한 수정 단계를 식별하고 올바른 궤적과 연결하여 실시간 복구

Preliminary

2.1 Task Formulation

term
- $π_θ$: LLM
- $τ_t$: trajectories = (a1, o1, ..., at, ot)
- $a, o$: action, observation. observation is obtained after executing action
- $a_{t+1}$: ~$π_θ(.|τ_t, u)$: trajectory와 query(u)로 부터 action을 생성함.
이 논문에서는 ReACT 방법론을 사용해서 행동 전에 근거를 생성함.
최대 라운드, 혹은 성공적으로 완료하면 종료됨.
최종 보상이 [0,1]로 주어짐.

전체 trajectory를 위와같이 수식화 할 수 있는데, 간단하게 - 매 action 생성시 greedy하게 action을 생성한다

2.2 MCTS

process
- 선택(Selection): UCT 전략으로 확장할 다음 노드 선택
- 확장(Expansion): 선택된 노드를 기반으로 새 노드 생성 및 트리에 추가
- 시뮬레이션(Simulation): 새 노드에서 최종 노드까지 다중 롤아웃 수행
- 역전파(Backpropagation): 시뮬레이션 결과 기반으로 노드 값 업데이트

3. Method

두가지로 구성
- Phase I: Model-Guided Reflection Trajectory Generation - MCTS사용해서 틀린 궤적을 수정된 궤적으로 변환
- Phase II: Iterative Self-Training with Revision Trajectories - 에이전트는 수정 궤적에 대해 반복적으로 훈련

3.1 Phase I: Model-Guided Reflection Trajectory Generation

Reflection Trajectory Definition
- Initial Trajectory: $τ^i = (a^i_1, o^i_1, ..., a^i_t, o^i_t)$
- Bad Trajectory: $τ^b = (τ^i, a^b_{t+1}, o^b_{t+1}, ..., a^b_{T_b}, o^b_{T_b})$ - 오류가 있거나 낮은 보상 궤정
- Good Trajectory: $τ^g = (τ^i, a^g_{t+1}, o^g_{t+1}, ..., a^g_{T_g}, o^g_{T_g})$ - 최적 또는 높은 보상이 있는 궤적
- Revision Trajectory: $τ^r = (τ^i, a^b_{t+1}, o^b_{t+1}, ..., a^b_{t'}, o^b_{t'}, rs, a^g_{t+1}, o^g_{t+1}, ..., a^g_{T_g}, o^g_{T_g})$ - 중간에 rs가 들어감.
  - rs: 전환을 표시하는 special prompt - 에이전트가 과거 행동에 대해 반영할 수 있도록 하는 Assistant: [reflection on the error] \n Human: OK.와 같은 간단한 통신 프롬프트
- 상호작용은 궤적이 종료된 후 최종 보상을 제공.
  - $r(τ^b) < β < r(τ^g) ≤ 1, α < r(τ^g) = r(τ^r)$
    - β: 나쁜 궤적과 좋은 궤적의 보상을 구분하는 임계값
    - α: 고품질 궤적의 하한선
      예) $r(τ^g) = 1 - 최적의 궤적을 의미함$
- 사실, 논문에서는 보상에 대한 부분을 정확하게 서술하지 않음.
  - 추측이나, 데이터셋에 따라 다르게 설정할 수 있을 것 같음. 정확도와 같은것으로 하거나, LLM as judge를 사용할 수도 있음.
  - 부록 A
Trajectory Collection with MCTS
- MCTS를 사용해 수정된 궤적을 수집. query u에 해당하는 초기 노드 $s_0$부터 선택,확장,시뮬레이션,역전파를 반복적 수행
- Figure 2와 같이 미래 행동을 롤아웃으로 샘플링함.
- Selection:
  - $UCT(s) = Q(s) + c_{uct} \sqrt(\frac{log N_p(s)}{N(s)})$
    - Q(s): 상태 s의 평균 보상
    - N(s): 상태 s를 방문한 횟수
    - $N_p(s)$: s의 부모 노드 총 방문횟수
    - $c_{uct}$: 탐색 활용 균형 제어하는 상수
  - 종료 조건에 도달하면 reward를 계산함. 다양한 궤적을 구하는데, 좋은 궤적 및 나쁜 궤적은 앞에 일정 부분의 궤적을 공유
- 수정 궤적: rs에 대해 10가지 다른 수정 생각을 수동으로 작성 - 10개 다른 프롬프트를 준비 - 임의로 하나가 샘플링되어 전환지점에 추가됨.
Transition Point Determination with Actor Model
- 논문에서는 길게 서술했으나, 그냥 간단하게 - LLM보고 어디 지점이 처음으로 실수를 하는지 식별하게 함.(Appendix A)
- 잘못된 행동 $a_t$를 식별하면, 전환지점은 t로 설정되고 뒤에 생긴 나쁜 지점은 좋은 지점의 뒷부분을 insert함.

3.2 Phase II: Iterative Self-Training with Revision Trajectories

목표: MCTS를 통해 수집된 자체 생성 수정 궤적을 사용하여 언어 에이전트를 훈련

수정 궤적으로만 훈련: 초기 올바른 궤적 식별하는 능력이 떨어짐, 콜드 스타트 문제가 있음.
해결책: 단계적 난이도 조절
- 초기 단계 - 궤적 혼합 전략
  - 훈련 중 수정 궤적 및 "좋은 궤적"(완전 최적은 아니지만 높은 보상을 받는 궤적)을 함께 사용
  - 낮은 초기 임계값: 초기에는 "좋은 궤적"으로 인정받기 위한 임계값(α)을 낮게 설정(논문에서 첫 훈련에 α = 0.5로 설정).
    -> 초기에 더 많은 다양한 성공 사례를 학습할 수 있음
- 중간 단계 - 점진적 품질 향상
  - 임계값 점진적 증가: 훈련이 진행됨에 따라 α 값을 점진적으로 증가(2차 반복에서 α = 0.7, 3차 반복에서 α = 1.0).
    -> 임계값이 높아짐에 따라 "좋은 궤적"의 품질 기준이 높아져, 점차 최적 궤적에 가까워짐.
  - 반복적 개선: 각 반복마다 현재 모델을 사용해 새로운 수정 궤적을 생성, 이를 다시 훈련 데이터로 활용
- 전체 과정 - 일반화 능력 강화
  - 범용 데이터셋 통합: 수집한 궤적 데이터셋(revision, good)과 일반 대화 데이터셋을 섞어서 훈련함(AgentTuning 전략)
  - 혼합 비율 η: 손실 함수에서 η는 에이전트 특화 데이터와 일반 데이터 간의 비율을 조절(논문에서 η = 0.2로 설정).
- 반복 과정 - 훈련된 모델로 다시 궤적 데이터를 생성, α 조절 등을 함.
Loss:
- 윗줄에 있는 텀: 궤적 데이터 - SFT - 좋은 궤적은 다 학습하고, 수정 궤적은 수정 이후의 궤적만 학습함.
  -> 수정 궤적의 앞부분은 틀린 정보를 포함할 수 있기 때문에 학습하지 않는 것으로 보임.
- 밑줄에 있는 텀: ShareGPT데이터 셋

4. Experiment

4.1 Interactive and Agentic Environments

4.2 Experiment Setting

데이터

- WebShop 300개, SciWorld 200개, TextCraft 200개 시뮬레이션으로 MCTS 수행

- 나쁜/좋은 궤적 구분 위한 임계값 β = 0.2 설정

- 반복적 훈련: 1차(α=0.5) → 2차(α=0.7) → 3차(α=1.0)로 점진적 향상

MCTS

- 8회 롤아웃 샘플링

- 깊이 20으로 설정

- 각 깊이에서 4개 후보 생성

- 탐색/활용 균형을 위한 UCT 계수 cuct = 0.25

모델:

- 주요 모델: Llama-3.1-8B-Instruction

- 비교 대상: GPT 계열, Claude 3, AgentLM, Agent-Flan, ETO 등

4.3 Main Result

4.4 Findings with analysis

수정 궤적의 우수성
- Agent-R의 수정 궤적이 최적 궤적만 사용하는 것보다 더 효과적
- 반복적 자체 훈련으로 모델 능력 점진적 향상
- 수정 궤적과 최적 궤적 혼합이 성능 크게 향상시킴
- 최적 궤적에 단순히 좋은 궤적만 추가하면 오히려 성능 저하 (노이즈 발생)

자체 반영 능력
- Agent-R로 훈련된 모델은 실패 궤적 주어질 때 효과적으로 수정 가능
- direct-revision 방법보다 높은 자체 수정 성능 보임
- SFT나 DPO로 훈련된 모델들은 기본 성능은 높으나 자체 반영 능력 부족
- GPT-4o는 테스트 세트에서는 성능이 낮으나 자체 반영 능력은 뛰어남

오류 식별 능력
- Agent-R 훈련 모델은 궤적 내 오류를 더 빠르게 인식
- 평균 수정 길이가 짧아짐 = 오류를 더 빨리 발견함
- 액터 모델의 향상된 반영 능력이 다른 모델 지원에도 활용 가능

회피 능력
- 최적 궤적만으로 훈련 시 반복적인 루프에 갇히는 문제 발생
- Agent-R 궤적으로 훈련 시 데드 루프 발생 크게 감소
- 수정 궤적 훈련 모델은 새로운 행동 탐색 능력과 루프 회피 능력 향상