STAIRS-Former: Spatio-Temporal Attention with Interleaved Recursive Structure Transformer for Offline Multi-task Multi-agent Reinforcement Learning

이 논문은 다양한 에이전트 수와 미시적 관측 환경에서 장기적 시간 의존성과 에이전트 간 협력을 효과적으로 포착하기 위해 공간 및 시간 계층 구조와 인터리브드 재귀적 구조를 도입한 STAIRS-Former 를 제안하고, 이를 통해 오프라인 다중 작업 다중 에이전트 강화학습에서 기존 방법들을 능가하는 새로운 최고 성능을 달성했음을 보여줍니다.

핵심

1. 문제 상황: "혼란스러운 축구 경기"

상상해 보세요. 11 명으로 구성된 축구 팀이 있는데, 경기 도중 갑자기 팀원이 1 명, 2 명, 혹은 3 명씩 사라지거나 새로 합류한다고 가정해 봅시다.

• 기존의 방법들 (ODIS, HiSSD 등): 이 팀들은 "지금까지 본 상황"을 단순히 나열해서 기억하려 했습니다. 하지만 팀원이 바뀌거나, 상대방이 갑자기 달라지면 "누가 누구를 지켜야 하지?", "어디로 공을 패스해야 하지?"를 판단하는 데서 혼란이 생겼습니다. 마치 한 줄로 서서 서로의 얼굴만 보고 있는 것처럼, 중요한 정보 (적의 위치, 아군의 상태) 를 골라내지 못하고 모든 정보를 다 똑같이 중요하게 여기느라 느려졌습니다.
• 결국: 새로운 상황 (예: 10 명이던 팀이 12 명이 됨) 에 가면 전혀 못 하거나, 아주 엉뚱한 실수를 반복했습니다.

2. 해결책: "STAIRS-Former" (계단식 지능)

이 연구팀은 **"STAIRS-Former"**라는 새로운 시스템을 만들었습니다. 이름처럼 **계단 (Stairs)**을 오르듯 정보를 정리하고, **시간 (Temporal)**과 **공간 (Spatial)**을 동시에 고려합니다.

이 시스템은 크게 세 가지 핵심 아이디어로 작동합니다.

① 공간적 계단 (Spatial Hierarchy): "현명한 코치"

기존 모델은 모든 선수 (에이전트) 와 모든 적을 똑같은 중요도로 보았습니다. 하지만 STAIRS-Former 는 코치처럼 행동합니다.

• 비유: 경기장에서 코치는 "지금 내 아군 1 번이 위험해! 저 적 2 번을 먼저 막아!"라고 외칩니다.
• 작동: 이 모델은 수많은 정보 중에서 **가장 중요한 것 (적의 위치, 아군의 부상 여부)**에 집중하고, 덜 중요한 것은 잠시 무시합니다. 마치 확대경으로 중요한 부분만 선명하게 보는 것과 같습니다.

② 시간적 계단 (Temporal Hierarchy): "단기 기억 vs 장기 기억"

반은 짧게, 반은 길게 기억하는 두 가지 메모리 시스템을 가지고 있습니다.

• 단기 기억 (Low-level): "지금 당장 적의 공이 날아왔어!" 같은 즉각적인 반응.
• 장기 기억 (High-level): "지난 10 분 동안 적들은 왼쪽으로 몰리는 경향이 있었어." 같은 흐름을 파악.
• 비유: 축구 선수가 현재의 공을 발로 차는 것 (단기) 과, 경기 전체의 흐름을 읽어서 다음 전략을 세우는 것 (장기) 을 동시에 잘합니다. 기존 모델은 이 두 가지를 섞어서 기억하느라 둘 다 제대로 못 했습니다.

③ 토큰 드롭아웃 (Token Dropout): "훈련용 가짜 부상"

이게 가장 재미있는 부분입니다. 훈련할 때 의도적으로 팀원 중 몇 명을 "실종" 시키거나 정보를 숨겨버립니다.

• 비유: 축구 연습 경기에서 코치가 "오늘은 1 번 선수가 다쳤으니, 10 명으로 경기를 해봐"라고 시킵니다.
• 효과: 이렇게 훈련을 받으면, 실제 경기에서 팀원이 갑자기 사라지거나 숫자가 달라져도 당황하지 않고 유연하게 대처할 수 있게 됩니다. 마치 다양한 상황에 대비한 근육을 키우는 것과 같습니다.

3. 결과: "새로운 세계 기록"

이 모델을 스타크래프트 (SMAC) 같은 복잡한 게임과 다양한 로봇 시뮬레이션에서 테스트했습니다.

• 기존 모델: 새로운 맵이나 팀원 수가 바뀌면 성적이 뚝 떨어졌습니다.
• STAIRS-Former: 어떤 상황에서도 압도적인 승리율을 기록했습니다. 특히, 훈련할 때 보지 못했던 새로운 팀 구성 (예: 3 대 3 에서 4 대 4 로 변함) 에도 완벽하게 적응했습니다.

4. 한 줄 요약

"STAIRS-Former 는 여러 대의 로봇이나 캐릭터가 협력할 때, '누가 중요한지'를 잘 골라내고 (공간), '과거의 흐름'을 잘 기억하며 (시간), '상황이 바뀌어도' 당황하지 않도록 (훈련 방식) 설계된, 마치 현명한 팀 코치 같은 인공지능입니다."

이 기술은 앞으로 자율 주행 차량들이 서로 소통하거나, 재난 구조용 드론들이 협력할 때 큰 도움이 될 것으로 기대됩니다.

기술 요약

1. 연구 배경 및 문제 정의 (Problem)

배경:
오프라인 다중 에이전트 강화학습 (Offline MARL) 은 실제 시스템 (연결된 차량, 협업 드론 등) 의 훈련 비용을 절감하고 안전성을 높이기 위해 유망한 접근법입니다. 특히, 다양한 작업 (Multi-Task) 과 다양한 에이전트 수를 가진 데이터셋에서 학습된 정책을 새로운 시나리오에 일반화하는 것이 핵심 과제입니다.

주요 문제점:
기존의 오프라인 다중 작업 MARL 방법론 (ODIS, HiSSD 등) 은 주로 UPDeT(Transformer 기반 확장 가능 아키텍처) 를 기반으로 합니다. 그러나 이러한 기존 방법들은 다음과 같은 한계를 가집니다:

1. Attention 메커니즘의 비효율적 활용: 에이전트 간의 협력 (Inter-agent coordination) 을 위한 Attention 메커니즘을 충분히 활용하지 못하며, 주로 관측 차원의 변동성을 처리하는 데만 집중합니다.
2. 단순한 히스토리 토큰: 단일 히스토리 토큰 (Single history token) 에 의존하여 부분 관측성 (Partially Observable) 환경에서 필요한 장기적인 시간적 의존성 (Long-horizon temporal dependencies) 을 포착하는 데 한계가 있습니다.
3. 심층적 관계 모델링 부재: 단일 레이어의 Transformer 를 사용하여 에이전트, 엔티티, 히스토리 간의 복잡한 관계를 심층적으로 모델링하지 못합니다.

2. 제안 방법론: STAIRS-Former

저자들은 이러한 한계를 극복하기 위해 STAIRS-Former를 제안합니다. 이는 공간적 (Spatial) 과 시간적 (Temporal) 위계를 통합한 Transformer 아키텍처로, 다음과 같은 세 가지 핵심 구성 요소로 이루어져 있습니다.

2.1. 공간 재귀 모듈 (Spatial Recursive Module)

• 목적: 국소 관측 내의 다양한 엔티티 (아군, 적군, 환경 객체) 간의 관계를 심층적으로 추론합니다.
• 구조: 'Spatial-Former'라고 불리는 재귀적 심층 Transformer 를 사용합니다.
  • 각 레이어는 공유 파라미터를 사용하여 $\nu$번 재귀적으로 적용됩니다.
  • 이전 레이어의 최종 상태와 현재 재귀 단계를 결합하여 ($z^{l}_{j+1} = f(z^{l}_j + z^{l-1})$) 더 깊은 관계 추론을 가능하게 합니다.
  • 이를 통해 에이전트가 가장 관련성 높은 엔티티에 집중할 수 있도록 Attention 을 유도합니다.

2.2. 시간 모듈 (Temporal Module)

• 목적: 부분 관측성 환경에서 단기 및 장기 의존성을 동시에 포착합니다.
• 구조: 두 가지 다른 주기로 업데이트되는 계층적 히스토리 상태를 도입합니다.
  • 저수준 히스토리 ($h^L$): 매 시간 단계마다 업데이트되며, 단기적인 세부 정보를 담습니다.
  • 고수준 히스토리 ($h^H$): $T_H$ 단계마다 GRU 를 통해 업데이트되며, 장기적인 맥락을 요약합니다.
• 특이점: Attention 블록 후 단일 MLP 를 사용하는 대신, 공간 엔티티 토큰과 히스토리 토큰을 위해 각각 독립적인 FFN(Fully Connected Network) 을 적용합니다. 이는 공간적 관계 추론과 시간적 추상화가 서로 간섭하지 않고 전문화되도록 합니다.

2.3. 토큰 드롭아웃 메커니즘 (Token-Dropout Mechanism)

• 목적: 에이전트 수나 엔티티 구성이 다른 미지의 작업 (Unseen tasks) 에 대한 일반화 성능을 향상시킵니다.
• 작동 방식: 학습 중 에이전트 토큰을 확률 $p_{drop}$로 무작위 드롭합니다.
  • 예외: 에이전트 자신의 토큰, 히스토리 토큰, 그리고 오프라인 정규화를 위해 데이터셋 행동과 연결된 토큰은 드롭하지 않습니다.
• 효과: 모델이 특정 엔티티 수에 과적합 (Overfitting) 되는 것을 방지하고, 다양한 에이전트 구성에 대한 강건성을 높입니다.

3. 주요 기여 (Key Contributions)

1. 새로운 Transformer 아키텍처: 오프라인 다중 작업 MARL 을 위한 STAIRS-Former를 제안하여, 토큰 간 중요한 정보를 선택적으로 집중 (Selective Attention) 하도록 설계했습니다.
2. 계층적 구조 도입: 공간적 재귀 (Spatial Recursion) 와 시간적 계층 (Temporal Hierarchy) 을 Transformer 내부에 통합하여, 다양한 에이전트 수와 부분 관측성 환경에서의 히스토리 의존성을 효과적으로 처리합니다.
3. 성능 입증: SMAC, SMAC-v2, MPE, MaMuJoCo 등 다양한 벤치마크에서 오프라인 다중 작업 데이터를 기반으로 실험을 수행하여, 기존 최첨단 (SOTA) 방법론 (HiSSD, ODIS 등) 을 압도하는 성능을 달성했습니다.

4. 실험 결과 (Results)

• 벤치마크: StarCraft Multi-Agent Challenge (SMAC, SMAC-v2), Multi-Agent Particle Environment (MPE), Multi-Agent MuJoCo (MaMuJoCo).
• 데이터셋: Expert, Medium, Medium-Expert, Medium-Replay 등 다양한 품질의 오프라인 데이터셋.
• 주요 성과:
  • SMAC (Marine-Hard, Stalker-Zealot): HiSSD 대비 평균 성능이 **39.5% ~ 48.6%**까지 향상되었습니다. 특히 불완전한 데이터 (Medium 등) 에서 큰 개선을 보였습니다.
  • 일반화 능력: 훈련 시 보지 못한 에이전트 수 (Unseen tasks) 에서도 높은 승률을 기록하며, 공간 - 시간적 추론과 토큰 드롭아웃의 시너지 효과를 입증했습니다.
  • SMAC-v2: 더 높은 확률적 요소와 복잡한 상호작용을 가진 환경에서도 기존 방법론을 크게 앞섰습니다.
  • Attention 분석: 시각화 결과, STAIRS-Former는 적군, 아군, 히스토리 토큰 등 상황에 따라 중요한 토큰에 집중하는 반면, 기존 방법은 토큰에 균일하게 Attention 을 분산시키는 것을 확인했습니다.

5. 의의 및 결론 (Significance)

이 논문은 오프라인 다중 에이전트 강화학습에서 Transformer 의 잠재력을 완전히 활용하지 못했던 기존 접근법의 한계를 지적하고, 이를 해결하는 새로운 아키텍처를 제시했습니다.

• 구조적 혁신: 단순히 Transformer 를 적용하는 것을 넘어, 공간적 재귀와 이중 시간 스케일을 도입하여 부분 관측성 환경에서의 장기적 의존성과 복잡한 에이전트 간 관계를 모델링하는 방법을 제시했습니다.
• 실용성: 에이전트 수가 변하는 동적 환경과 다양한 작업에 대한 일반화 능력을 크게 향상시켜, 실제 로봇 공학이나 자율 주행 등 복잡한 다중 에이전트 시스템에 적용 가능한 강력한 오프라인 학습 프레임워크를 제공합니다.
• 효율성: HiSSD 와 같은 기존 SOTA 모델보다 파라미터 수는 적으면서도 더 높은 성능을 달성하여, 효율적인 모델 설계의 중요성을 강조합니다.

결론적으로, STAIRS-Former는 오프라인 MARL 의 일반화 문제를 해결하기 위해 **구조화된 주의 (Structured Attention)**와 강건한 정규화가 필수적임을 보여주며, 해당 분야의 새로운 기준 (SOTA) 을 설정했습니다.