Efficient Reinforcement Learning for Long-Horizon Tool-Use Agentic Tasks

당신이 아주 복잡하고 며칠씩 걸리는 보드게임을 하는 똑똑한 로봇에게 가르치고 있다고 상상해 보세요. 이 게임은 단순히 주사위를 굴려 말을 움직이는 게임이 아닙니다. 로봇은 다른 플레이어들과 대화해야 하고, 특수한 도구 상자를 사용해야 하며, 엄격한 규칙을 따라야 하고, 자신이 실제로 이겼는지 알기 위해 며칠을 기다려야 합니다. 이것이 바로 '에이전틱 AI(Agentic AI)'의 세계입니다. 컴퓨터 프로그램이 디지털 세상에서 행동을 취하는 작은 에이전트처럼 행동하는 것이죠. 가장 큰 문제는 이 게임들이 너무 길고 복잡해서 로봇의 뇌(메모리)가 넘쳐버린다는 것입니다. 이는 마치 10시간 동안의 대화 속 모든 단어를 기억하면서 동시에 수학 문제를 푸는 것과 같습니다. 결국, 생각할 공간이 부족해지게 됩니다. 과학자들은 로봇이 정보의 무게에 짓눌려 뇌가 녹아내리지 않으면서도, 이러한 긴 게임을 더 잘 수행할 수 있도록 훈련하는 방법을 연구하고 있습니다.

이 논문은 SINKFLEX-RL이라는 새로운 훈련 시스템을 소개합니다. 이는 로봇이 더 큰 뇌를 갖지 않고도 게임의 가장 중요한 부분을 기억할 수 있게 도와주는 영리한 트릭과 같습니다. 연구진은 표준 게임 인터페이스, 실수로부터 배우는 스마트한 방법(이를 "그룹 상대적 정책 최적화" 또는 "GRPO"라고 부릅니다), 그리고 로봇의 주의력(attention)을 위한 특별한 메모리 절약 기술을 결합한 시스템을 구축했습니다. 과거 8,000단계의 대화 내용을 하나하나 다 기억하려고 노력하는 대신, 이 시스템은 "싱크(sink)" 메커니즘을 사용합니다. 이것은 로봇의 마음속에 있는 마법의 스펀지라고 생각하면 됩니다. 이 스펀지는 오래된 정보의 압도적인 노이즈를 흡수하지만, 로봇이 안정성을 유지하는 데 도움이 되는 필수적인 "싱크"(시작점)는 남겨둡니다. 이 스펀지를 사용함으로써, 로봇은 과거의 무게에 짓눌리지 않고 현재의 움직임에 집중할 수 있습니다.

연구팀은 로봇이 고객을 돕고, 재고를 확인하기 위해 도구를 사용하며, 매장 정책을 따르는 시뮬레이션된 소매점 환경에서 이 시스템을 테스트했습니다. 초기 테스트에서 로봇의 성능 점수는 학습함에 따라 0.25에서 0.44로 상승했습니다. 하지만 진짜 마법은 시스템이 얼마나 많은 컴퓨터 메모리(VRAM)를 필요로 하는지 테스트했을 때 일어났습니다. 대화가 매우 길어질 때(8,192 토큰), 기존의 표준적인 사고 방식은 메모리가 부족하여 충돌(crash)이 발생했습니다. 그러나 새로운 SINKFLEX-RL 시스템은 25.53 GB의 메모리만을 사용하며 계속 작동했습니다. 심지어 길이가 약간 짧은 4,096 토큰에서도, 새 시스템은 기존 방식에 비해 무려 19.7%의 메모리를 절약했습니다. 저자들은 이것이 매우 비싼 하드웨어 없이도 이러한 장기적(long-horizon) 에이전트를 훈련하는 것이 가능하다는 것을 증명한다고 제안하지만, 이것은 예비적인 검토일 뿐 최종적이고 완벽한 해결책은 아니라고 언급했습니다.

문제점: 로봇의 메모리 누수

당신이 이 이야기 속의 로봇이라고 상상해 보세요. 당신은 고객이 셔츠를 사는 것을 돕는 게임을 하고 있습니다. 당신은 고객에게 사이즈를 묻고, 고객은 대답하고, 당신은 재고를 확인하고, 고객은 다른 색상을 요구하고, 당신은 다시 확인하는 식입니다. 50번의 턴이 지난 후, 당신은 고객이 처음에 말했던 것을 기억하여 사이즈를 잊지 않았는지 확인해야 합니다.

AI의 세계에서는 이를 "장기적(long-horizon)" 과업이라고 부릅니다. 문제는 대화가 길어질수록 컴퓨터가 그 모든 단어를 보유하는 데 필요한 메모리 양이 엄청나게 빠르게 증가한다는 것입니다. 이는 발걸음을 옮길 때마다 점점 더 무거워지는 배낭을 메고 가는 것과 같습니다. 만약 대화가 너무 길어지면(예: 8,000 단어), 배낭이 너무 무거워져서 로봇이 쓰러지게 됩니다. 이를 VRAM(비디오 램)이 부족하다고 합니다. VRAM은 컴퓨터가 사고하는 데 사용하는 고속 메모리입니다.

연구진은 표준적인 로봇 훈련 방식이 "메모리 벽"에 부딪히고 있다는 점을 발견했습니다. 로봇들은 모든 것을 명시적으로 기억하려고 했기 때문에 너무 많은 비용이 들었고, 이로 인해 정체되었습니다. 그들은 추론 능력을 잃지 않으면서도 효율적인 방법이 필요했습니다.

해결책: 마법의 스피지와 팀 허들

논문은 이 메모리 누수를 해결하기 위해 SINKFLEX-RL이라 불리는 세 부분으로 구성된 솔루션을 제안합니다.

1. Gymnasium Wrapper (유니버설 리모컨)
먼저, 그들은 로봇이 다양한 게임 환경과 대화할 수 있게 해주는 유니버설 리모컨 같은 표준 인터페이스를 구축했습니다. 로봇이 상점에 있든, 은행에 있든, 여행사에 있든, 이 래퍼(wrapper)는 로봇이 도움을 요청하고, 도구를 사용하고, 피드백을 받는 방법을 확실히 알 수 있게 해줍니다. 이는 로봇에게 표준적인 규칙 세트를 제공하여, 매 게임을 할 때마다 새로운 언어를 배울 필요가 없도록 만드는 것과 같습니다.

2. 그룹 허들 (GRPO)
다음으로, 그들은 로봇이 학습하는 방식을 변경했습니다. 보통 로봇이 학습하려면, 어떤 움직임이 좋았는지 알려주는 별도의 "코치"(가치 모델)가 필요합니다. 하지만 이 코치는 추가적인 메모리를 차지합니다. 대신, 이 시스템은 **그룹 상대적 정책 최적화(GRPO)**라고 불리는 방법을 사용합니다.

상상해 보세요. 로봇이 똑같은 게임을 10번 연속으로 플레이하는데, 매번 선택이 조금씩 다릅니다. 코치에게 묻는 대신, 로봇은 자기 자신의 10가지 버전을 살펴봅니다. 로봇은 이렇게 말합니다. "헤이, 3번 버전이 1번 버전보다 점수가 높았으니, 3번의 움직임을 따라 해야지." 로봇은 무엇이 효과적이었는지 알아내기 위해 자기 자신의 그룹과 비교합니다. 이는 학생들이 각자의 답을 채점하기 위해 선생님을 기다리는 대신, 서로의 답을 비교하며 누가 맞혔는지 확인하는 스터디 그룹과 같습니다. 이 방식은 별도의 코치를 훈련시킬 필요가 없기 때문에 엄청난 양의 메모리를 절약합니다.

3. 마법의 스피지 (Sink-Aware FlexAttention)
이것이 가장 창의적인 부분입니다. 로봇의 "주의력(attention)"은 대화 중 어떤 단어가 중요한지 결정하는 방식입니다. 긴 대화에서 로봇은 보통 모든 단어를 보려고 시도하는데, 이는 느리고 메모리를 많이 사용합니다.

연구진은 긴 대화에서 대화의 시작 부분이 "싱크(sink)" 역할을 한다는 점을 깨달았습니다. 즉, 로봇의 주의력이 안정성을 유지하기 위해 자연스럽게 모이는 곳입니다. 그들은 FlexAttention이라는 수학적 기법을 사용하여 로봇이 이러한 "싱크" 단어들을 다르게 처리할 수 있게 했습니다.

이렇게 생각해 보세요. 만약 당신이 100페이지짜리 책을 읽고 있다면, 책 전체를 한꺼번에 손에 들고 있을 필요는 없습니다. 첫 페이지(싱크)와 지금 읽고 있는 페이지를 잡고, 중간 페이지들은 필요할 때까지 사라지게 둘 수 있습니다. "싱크"는 당신이 전체 책을 다 짊어지지 않고도 이야기의 맥락을 유지할 수 있게 해주는 북마크와 같습니다. 이 시스템은 "제로 밸류 싱크(zero-value sink)"를 사용하는데, 이는 수학적으로 "우리는 오래된 단어들의 실제 데이터를 저장할 필요는 없지만, 균형을 잡기 위해 그 단어들이 그곳에 있었다는 사실은 알아야 한다"는 것을 의미합니다. 이를 통해 로봇은 메모리 부족 없이 긴 대화(최대 8,192 토큰)를 처리할 수 있습니다.

결과: 효과가 있는가?

팀은 시뮬레이션된 소매점 환경에서 이 새로운 시스템을 테스트했습니다. 그들은 로봇이 일정 기간 동안 학습하는 과정을 관찰했습니다.

• 학습 진행: 훈련 시작 시 로봇의 고객 응대 점수는 0.25였습니다. 관찰된 훈련 기간이 끝날 무렵, 점수는 0.44로 상승했습니다. 연구진은 또한 "훈련 점수"와 "궤적 보상(trajectory reward)"(로봇이 얼마나 잘하고 있는지 나타내는 내부 체크마크와 같은 것)이 각각 0.18에서 0.40, 그리고 0.39에서 0.39로 상승하는 것을 확인했습니다. 이는 로봇이 실제로 학습하고 개선되고 있음을 시사하지만, 저자들은 이것이 이 방법이 완벽하다는 최종적인 증명이 아니라 예비적인 검토라는 점을 주의 깊게 언급했습니다.

• 메모리 절감: 가장 흥미로운 결과는 메모리에 관한 것이었습니다. 그들은 다양한 대화 길이로 시스템을 테스트했습니다.

  • 4,096 토크에서 기존 방식은 28.06 GB의 메모리가 필요했습니다. 새로운 SINKFLEX-RL 시스템은 22.52 GB만을 필요로 했습니다. 이는 5.54 GB, 즉 **19.7%**의 절감 효과를 가져왔습니다.
  • 8,192 토큰에서 기존 방식은 메모리 부족으로 완전히 충돌(OOM)했습니다. 그러나 새로운 시스템은 계속 실행되었으며, 단 25.53 GB의 메모리만을 사용했습니다.

이것이 의미하는 바 (그리고 그렇지 않은 것)

이 논문은 표준 게임 인터페이스, 스마트한 그룹 학습 방법, 그리고 메모리 절약형 어텐션 기법을 결합함으로써, 슈퍼컴퓨터 없이도 매우 길고 복잡한 작업을 수행하는 로봇을 훈련시키는 것이 가능하다는 것을 보여줍니다. "싱크" 기법은 스펀지처럼 작용하여 메모리 압력을 흡수함으로써 로봇이 계속 나아갈 수 있게 합니다.

하지만 저자들은 자신들이 아직 증명하지 못한 부분에 대해서도 매우 정직하게 밝히고 있습니다. 그들은 이것이 학습하는 가장 좋은 방법인지, 혹은 모든 종류의 로봇에 작동하는지를 테스트하지 않았습니다. 그들은 로봇이 얼마나 더 빨리 학습하는지도 측정하지 않았으며, 단지 충돌 없이 학습할 수 있다는 것만을 보여주었습니다. 또한, 이것이 "개념 증명(proof of concept)", 즉 아이디어가 작동함을 보여주는 성공적인 테스트 실행이지, 완성된 제품은 아니라고 언급했습니다.

요약하자면, SINKFLEX-RL은 AI 에이전트가 뇌가 폭발하지 않고도 긴 대화를 기억할 수 있도록 돕는 유망한 새로운 도구입니다. 이는 우리가 적절한 트릭을 사용한다면, 메모리를 관리할 수 있는 방법을 제공함으로써 복잡한 다일(multi-day) 과업을 수행할 수 있는 똑똑한 로봇을 만들 수 있음을 시사합니다.