대규모 언어 모델을 활용한 강화학습 인터페이스 발견

강화학습에서 에이전트가 무엇을 보고 어떤 보상을 받을지 설계하는 인터페이스 엔지니어링은 수동 작업이 많이 필요한 병목 구간이다. 이 논문은 LLM을 활용해 관측값 매핑과 보상 함수를 파이썬 코드로 자동 생성하고 진화시켜 사람이 설계한 것보다 효율적인 인터페이스를 찾아낸다.

강화학습 에이전트는 환경의 원시 상태(Raw State)를 그대로 보기보다 학습에 유리하게 가공된 관측값(Observation)을 입력받을 때 더 잘 학습한다. 예를 들어 로봇의 관절 각도 데이터만 주는 것보다 목표물과의 상대적 거리나 방향 벡터를 계산해서 주는 것이 신경망이 관계를 파악하는 데 훨씬 유리하다. LIMEN은 이러한 가공 과정을 LLM이 파이썬 코드로 직접 작성하게 한다.

기존의 자동 보상 설계 연구들은 관측값은 고정되어 있다고 가정했으나, LIMEN은 관측값 매핑 함수와 보상 함수를 하나의 쌍으로 묶어 진화시킨다. 이는 마치 생명체가 특정 먹이를 먹기 위해 시각 체계와 보상 체계를 동시에 진화시키는 것과 유사하다. LLM은 이전 세대의 코드와 그 코드로 학습한 에이전트의 성공률 피드백을 받아 코드를 수정하며, 이 과정에서 MAP-Elites 알고리즘을 통해 관측값의 복잡도와 보상의 구조가 다양한 여러 후보군을 유지한다.

결과적으로 LIMEN은 사람이 수동으로 설계하기 어려운 복잡한 기하학적 특징 추출이나 단계별 보상 설계를 스스로 수행한다. 이는 에이전트가 해결해야 할 문제의 난이도를 낮추어 학습 효율을 극대화하는 효과를 가져온다.

LIMEN은 상위 루프에서 인터페이스를 진화시키고 하위 루프에서 정책을 학습하는 이단계 최적화(Bilevel Optimization) 구조를 가진다. 상위 루프에서는 LLM이 관측값 매핑 함수 phi와 보상 함수 R을 포함하는 파이썬 프로그램을 생성한다. 생성된 프로그램은 구문 검사 및 실행 테스트를 거쳐 유효성이 검증된다.

하위 루프에서는 생성된 인터페이스를 사용하여 PPO 알고리즘으로 에이전트를 학습시킨다. 이때 학습된 정책의 에피소드 성공률을 계산하여 이를 인터페이스의 적합도(Fitness) 점수로 사용한다. [성공률 0~1 사이의 값 → 적합도 점수로 변환 → 아카이브 업데이트] 과정을 통해 우수한 인터페이스가 선택된다.

다양성 유지를 위해 MAP-Elites 아카이브를 사용하며, 관측값의 차원 수와 보상 함수의 구조적 복잡도(AST 노드 수)를 기준으로 2차원 그리드에 후보군을 배치한다. LLM은 아카이브에서 부모 인터페이스를 선택하고, 작업 설명과 이전 실패 사례의 오류 추적(Error Trace) 정보를 포함한 프롬프트를 입력받아 코드를 변이(Mutation)시킨다.

XLand-MiniGrid의 세 가지 그리드월드 태스크와 MuJoCo 기반의 로보틱스 태스크(Go1 Push Recovery, Panda Tracking)에서 실험을 진행했다. LIMEN은 모든 태스크에서 사람이 설계한 기본 인터페이스나 보상만 최적화한 기존 방식보다 높은 성공률을 기록했다. 특히 Hard 그리드월드 태스크에서 LIMEN은 85%의 성공률을 달성한 반면, 보상만 최적화한 방식은 1% 미만의 성공률로 실패했다.

Ablation 연구 결과, 관측값만 최적화하거나 보상만 최적화할 경우 특정 도메인에서 치명적인 성능 저하가 발생함을 확인했다. 예를 들어 Panda Tracking 태스크에서 관측값만 진화시켰을 때는 0%의 성공률을 보였으나, 보상과 함께 진화시켰을 때는 67%까지 성능이 향상되었다. 이는 관측값과 보상의 공동 설계(Co-design)가 필수적임을 시사한다.

또한 진화된 인터페이스는 물리적 수치가 변하는 환경 변화(Distribution Shift) 상황에서도 성능이 급격히 붕괴되지 않고 유지되는 강건함을 보였다. LLM이 단순히 수치를 맞추는 것이 아니라 태스크 해결에 필요한 구조적 특징을 추출하는 코드를 생성했기 때문이다.

LIMEN의 핵심은 인터페이스를 실행 가능한 프로그램(Executable Program)으로 취급하여 탐색 공간을 정의한 것이다. 관측값 매핑 함수 phi는 S -> O (최대 512차원 벡터)로, 보상 함수 R은 S x A x S -> R로 정의된다. 모든 코드는 JAX와 호환되도록 작성되어 GPU 상에서 대규모 병렬 학습이 가능하다.

진화 전략으로는 MAP-Elites를 사용하여 품질-다양성(Quality-Diversity)을 확보한다. 아카이브의 행동 기술자(Behavioral Descriptors)는 관측값의 차원 수와 보상 함수의 추상 구문 트리(AST) 노드 수로 설정되어, 단순한 인터페이스부터 복잡한 인터페이스까지 폭넓게 탐색한다. LLM은 Claude Sonnet 4.6을 사용하며, 온도 0.7 설정으로 창의적인 변이를 유도한다.

프롬프트 엔지니어링 측면에서는 'Evaluation Cascade' 기법을 도입했다. 모든 후보를 풀 트레이닝하는 대신 짧은 예비 학습을 통해 최소 기준을 통과한 후보만 정밀 평가하여 계산 비용을 절감한다. 또한 LLM에 제공되는 피드백에는 성공률뿐만 아니라 실패한 코드의 런타임 에러 로그를 포함하여 코드의 견고성을 높였다.

LIMEN은 신뢰할 수 있는 궤적 수준의 성공 메트릭(Success Metric)이 존재해야 진화가 가능하다는 한계가 있다. 또한 JAX 기반의 병렬 시뮬레이션 환경에서는 비용이 효율적이지만, 시뮬레이션 속도가 느린 환경에서는 인터페이스 평가를 위한 계산 비용이 병목이 될 수 있다.