LaGO: Latent Action Guidance for Online Reinforcement Learning

기술 요약: LaGO – 온라인 강화 학습을 위한 잠재 행동 가이드 (Latent Action Guidance for Online Reinforcement Learning)

문제 정의
대규모 언어 모델(LLM)은 계획 수립 및 순차적 의사 결정 분야에서 상당한 잠재력을 보여주었으나, 기존 방식들은 LLM을 직접적인 제어기(controller), 명시적인 고수준 플래너(high-level planner), 또는 세계 모델 시뮬레이터(world model simulator)로 사용하는 데 의존하는 경우가 많습니다. 이러한 방법들은 LLM이 정밀한 행동, 상세한 계획, 또는 정확한 미래 궤적을 생성할 것을 요구합니다. 그러나 이러한 요구사항은 실제 적용 시 지나치게 강력한 제약이 될 수 있습니다. LLM의 능력은 매우 다양하며, 작은 예측 오류라도 시간이 흐름에 따라 누적되어 정밀한 실행과 장기적 추론을 요구하는 강화 학습(RL) 환경에서 저조한 제어 성능을 초래할 수 있기 때문입니다. 결과적으로, 특히 관측값과 행동이 자연스럽게 언어나 시각적 형태로 표현되지 않는 벡터화된 환경(vectorized environments)에서는 LLM을 주요 제어기로 직접 사용하는 것이 신뢰하기 어렵습니다.

방법론
이러한 과제를 해결하기 위해, 저자들은 사전 학습된 LLM을 명시적인 제어기가 아닌, 잠재 행동 사전 지식(latent action priors)을 통한 "소프트(soft)" 가이드의 원천으로 활용하는 2단계 프레임워크인 **LaGO (Latent Action Guidance for Online Reinforcement Learning)**를 제안합니다.

1. 1단계: 오프라인 사전 지식 학습 (Offline Prior Learning)
   프레임워크는 오프라인 전문가 데이터셋($D$)으로부터 잠재 정책 사전 지식(latent policy prior)을 학습하는 것으로 시작합니다. 사전 학습된 LLM 백본($F_\phi$)을 사용하며, 데모로부터 추출된 비텍스트 형태의 상태-행동 시퀀스는 학습 가능한 투영 모듈($E_s$ 및 $E_a$)을 통해 잠재 표현(latent representations)으로 변환됩니다. 이 임베딩들은 텍스트 기반 작업 설명($x$)과 결합되어 LLM에 입력됩니다. 새롭게 추가된 액션 헤드($G_\theta$)는 궤적 접두사($\tau_t$)와 작업 설명이 주어졌을 때 다음 행동($a_{t+1}$)을 예측하도록 지도 학습 기반 미세 조정(supervised fine-tuning)을 통해 학습됩니다. 목적 함수는 전문가 데이터의 음의 로그 가능도(negative log-likelihood)를 최소화하는 것이며, 이를 통해 최적의 제어기보다는 거친 행동 사전 지식(coarse behavioral prior)을 포착하는 잠재 정책 사전 지식 $\pi_K(a_{t+1} | x, \tau_t)$을 생성합니다.

2. 2단계: 잠재 가이드를 이용한 온라인 RL (Online RL with Latent Guidance)
   두 번째 단계에서는 학습된 잠재 사전 지식 $\pi_K$를 동결(freeze)하고, 이를 온라인 RL 정책 $\pi_\psi(a|s)$ (예: PPO 또는 SAC)를 가이드하는 데 사용합니다. RL 에이전트는 환경과 상호작용하며 롤아웃(rollout)을 수집하지만, 정책 최적화는 잠재 사전 지식을 향해 정규화(regularized)됩니다. 전체 학습 손실은 다음과 같이 정의됩니다:
   $$L_{online} = L_{PG} + \beta L_{prior}$$
   여기서 $L_{PG}$는 표준 정책 경사(policy gradient) 목적 함수이며, $\beta$는 사전 지식의 영향력을 조절합니다.

   • **이산 행동 공간(discrete action spaces)**의 경우, 사전 지식 손실은 사전 지식이 예측한 행동의 음의 로그 가능도입니다.
   • **연속 행동 공간(continuous action spaces)**의 경우, 사전 지식 손실은 현재 정책과 사전 분포 사이의 교차 엔트로피(cross-entropy) 항으로 구현되며, 이는 (고정된 사전 지식의 엔트로피까지 고려했을 때) KL-발산 정규화와 동일한 최적화 신호 역할을 합니다.

이러한 설계는 RL 정책이 상호작용 중에 정확한 결정을 내릴 필요 없이, LLM에 인코딩된 구조적 지식으로부터 도움을 받을 수 있게 합니다.

주요 기여

• 새로운 프레임워크: LaGO는 사전 학습된 LLM의 지식을 온라인 RL을 위한 잠재 행동 사전 지식으로 활용하는 방법을 도입하여, LLM을 직접적인 제어기로 사용할 때 발생하는 엄격한 정확도 요구사항을 완화합니다.
• 2단계 아키텍처: 오프라인 사전 지식 학습(전문가 데이터에 대한 지도 학습 기반 미세 조정)과 온라인 정책 최적화(정규화를 통한 학습)를 분리함으로써, 사전 지식이 단독 에이전트로 작동할 만큼 충분히 정확하지 않더라도 시스템이 LLM의 지식을 활용할 수 있도록 합니다.
• 일반화: 프레임워크는 특정 행동 공간에 맞게 사전 지식 정규화 항을 조정함으로써 이산 및 연속 제어 도메인 모두에서 작동하도록 설계되었습니다.

실험 결과
저자들은 CLEVR-Robot(이산 제어)과 Meta-World(연속 제어)라는 두 가지 벤치마크에서 LaGO를 평가했습니다. 실험에서는 네 가지 작업 카테고리(Real, Rephrase, Easy, Hard)에 대해 LaGO를 "Vanilla PPO"와 비교했습니다.

• 성능 향상: LaGO는 두 환경 모두에서 보상(reward)과 성공률(success rate) 측면에서 Vanilla PPO보다 일관되게 우수한 성능을 보였습니다.
  • CLEVR-Robot에서 평균 성공률이 15.1%에서 27.2%로 증가했습니다.
  • Meta-World에서 평균 성공률이 2.7%에서 15.2%로 증가했습니다.
• 일반화: "Easy" 및 "Hard" 미학습 작업(unseen tasks)을 포함한 모든 작업 카테고리에서 개선이 관찰되었으며, 이는 사전 지식이 단순한 암기가 아닌 일반화 가능한 가이드를 제공함을 시사합니다.
• LLM 품질의 영향: Meta-World에서 Llama-2-7B 기반 사전 지식과 TinyLlama-1.1B 기반 사전 지식을 비교한 어블레이션 연구(ablation study) 결과, 더 강력한 사전 학습 모델이 더 효과적인 잠재 사전 지식을 생성함을 입증했습니다. Llama-2-7B를 사용한 LaGO가 일관되게 가장 높은 성능을 달-성한 반면, TinyLlama 사전 지식은 때때로 Vanilla PPO보다 낮은 성능을 보였는데, 이는 파운데이션 모델의 품질이 중요함을 강조합니다.

의의 및 주장
본 논문은 LaGO가 사전 학습된 LLM의 지식을 강화 학습에 통합하는 단순하고 실용적인 프레임워크를 제공한다고 주장합니다. 주요 의의는 LLM이 직접적인 제어기로 기능할 만큼 정확하지 않더라도, 온라인 의사 결정 및 계획 수립 능력을 향상시킬 수 있음을 입증했다는 점에 있습니다. LLM을 명시적인 플래너가 아닌 부드러운(soft) 잠재 가이드의 원천으로 취급함으로써, 프레임워크는 오류 누적의 위험을 완화하는 동시에 사전 학습 과정에서 습득한 풍부한 귀납적 편향(inductive biases)과 세계 지식을 활용합니다.

저자들은 사전 학습된 LLM에 인코딩된 지식과 일반화 능력이 계획 및 RL을 위한 강력한 사전 지식 역할을 할 수 있다고 결론짓습니다. 또한, 파운데이션 모델이 발전함에 따라 이러한 가이드의 효과도 자연스럽게 증가할 것이며, 이는 상호작용 중 정밀한 행동 생성을 요구하지 않고도 복잡한 의사 결정 과제에서 LLM을 활용할 수 있는 확장 가능한 경로를 제시한다고 제언합니다.