의료 에이전트를 위한 HEALTHCARE AI GYM

기존 의료 AI는 단발성 질문 답변에 치중되어 실제 진료 현장의 다단계 의사결정 과정을 반영하지 못했다. 이 논문은 135개의 전문 도구와 3,600개 이상의 작업을 포함한 강화학습 환경을 구축하여, 의료 AI가 실제 의사처럼 도구를 사용하고 추론하며 최적의 치료 경로를 찾을 수 있는 기반을 마련했다.

기존의 대형 언어 모델 학습은 주로 한 번의 질문에 한 번의 답변을 내놓는 방식에 최적화되어 있다. 하지만 실제 의료 현장은 환자의 상태를 묻고, 검사를 처방하고, 결과를 해석하는 다단계(Multi-turn) 상호작용이 필수적이다. 강화학습을 이 과정에 적용할 때, 최종 결과에만 보상을 주는 방식은 모델이 중간 단계에서 어떤 도구를 써야 할지 갈피를 못 잡게 만들며, 결국 도구 사용을 포기하고 긴 설명만 늘어놓는 '독백' 상태로 빠지게 한다.

이 논문은 이를 해결하기 위해 '선생님 모델(Teacher)'이 학생 모델의 학습 과정을 매 턴마다 세밀하게 지도하는 방식을 도입했다. 단순히 정답만 알려주는 것이 아니라, 정답에 이르는 과정에서 모델이 내린 결정(도구 사용 등)이 적절했는지를 KL Divergence라는 지표를 통해 실시간으로 교정한다. 이때 선생님 모델은 학생 모델의 과거 가중치를 서서히 반영하는 EMA(Exponential Moving Average) 방식을 사용하여 학습의 급격한 변화를 막고 안정성을 확보한다.

결과적으로 모델은 불필요하게 긴 답변을 내놓는 대신, 필요한 시점에 적절한 의료 도구를 호출하고 논리적인 추론 단계를 밟아가는 능력을 갖추게 된다. 이는 복잡한 의료 데이터를 다루는 에이전트가 단순히 지식을 암기하는 수준을 넘어, 실제 진료 프로세스를 수행하는 지능을 갖게 됨을 의미한다.

HEALTHCARE AI GYM은 Gymnasium 인터페이스를 기반으로 설계되어 step(action)과 render() 함수를 통해 표준화된 강화학습 환경을 제공한다. 135개의 도구는 증거 검색, 임상 평가, 중재 조치, 추론 보조 도구로 분류되며, OpenAI 호환 함수 호출 스키마를 따른다. 보상 함수는 Accuracy, Process Quality, Safety, Format, Coherence의 5가지 차원을 가중 합산하여 계산한다.

TT-OPD(Turn-Level Truncated On-Policy Distillation)는 학생 모델 θS와 EMA 기반의 선생님 모델 θT를 활용한다. 선생님 모델은 결과에 대한 특권 정보(Outcome-privileged information)를 입력받아 각 턴 t에서의 행동 a에 대한 로그 확률 πθT(at|st+)를 계산한다. [학생 모델의 정책 분포와 선생님 모델의 정책 분포를 입력으로] → [각 턴마다 KL Divergence를 계산하여] → [전체 궤적에 대한 손실 함수 LTT-OPD를 산출하고] → [이를 GRPO 손실 함수와 합산하여 가중치를 갱신한다].

학습 안정성을 위해 선생님 모델은 θT ← αθT + (1-α)θS (α=0.995) 식을 통해 5단계마다 업데이트되며, 30단계마다 학생 모델의 가중치를 강제로 복사하는 하드 카피 폴백을 수행한다. 또한 응답 길이 폭발을 막기 위해 코사인 함수 기반의 길이 제어 보상(Cosine length-controlled reward)을 적용하여, 응답이 최대 길이에 가까워질수록 보상을 삭감한다.

TT-OPD는 18개 벤치마크 중 10개에서 최고 성능을 기록했다. 특히 MedQA(USMLE)에서 87.1%를 달성하여 강화학습 전 베이스라인 대비 16.4%p 향상된 결과를 보였다. MIMIC-III(62.7%)와 eICU(57.1%) 등 실제 전자의무기록(EHR) 기반 추론 과제에서도 기존 GRPO보다 우수한 성능을 입증했다.

학습 역학 분석 결과, 일반 GRPO는 응답 길이가 7.7K에서 10.8K 토큰 사이에서 심하게 진동하며 불안정한 모습을 보인 반면, TT-OPD는 5.7K에서 9.3K 사이로 길이를 안정적으로 제어했다. 또한 에피소드당 평균 턴 수가 일반 모델은 단일 턴으로 급격히 감소(Collapse)하는 경향을 보였으나, TT-OPD는 7.0~7.4턴을 꾸준히 유지하며 다단계 도구 사용 능력을 보존했다.

본 연구는 다회차 에이전트 강화학습에서 발생하는 '에이전트적 붕괴(Agentic Collapse)'를 해결하기 위해 TT-OPD를 제안한다. 이는 기존의 OPSD나 SRPO가 단일 회차 추론에 집중했던 것과 달리, 전체 대화 궤적에 걸쳐 밀집된(Dense) 보상 신호를 제공한다. 특히 '결과 조건부 특권 힌트(Outcome-conditioned privileged hints)'를 선생님 모델의 컨텍스트에만 주입하여, 학생 모델이 직접 정답을 보지 않고도 정답에 가까운 추론 분포를 모방하도록 유도한다.

수학적으로는 Fisher Information Matrix F(θS)를 기반으로 한 KL penalty gradient가 복원력(Restoring force)으로 작용하여, 정책이 급격하게 변할 때 학습률을 암시적으로 조절하는 효과를 낸다. 또한 5D 보상 함수에서 발생하는 'Gradient Signal Dilution' 문제를 분석하여, 형식(Format) 보상이 정확도(Accuracy) 보상의 신호를 희석시키는 현상을 TT-OPD의 증류 손실이 보완함을 이론적으로 뒷받침했다.

현재의 보상 체계는 최종 결과에 의존하는 희소 보상(Sparse reward) 성격이 강해, 아주 긴 에피소드에서의 신용 할당(Credit assignment) 문제가 여전히 존재한다. 또한 9B 규모의 모델에서 실험되었으므로 더 큰 모델이나 20턴 이상의 초장기 에피소드에서의 확장성은 추가 검증이 필요하다.