💻 computer science

Scaling Short-Term Memory of Visuomotor Policies for Long-Horizon Tasks

이 논문은 장기 과업을 위한 시각 운동 정책(visuomotor policies)에서 단기 기억을 효과적으로 확장하는 게이트 어텐션(gated attention)과 계층적 압축(hierarchical compression)을 갖춘 트랜스포머 기반 아키텍처인 PRISM을 소개하며, 체계적인 평가를 위한 ReMemBench 벤치마크와 함께 대규모 사전 학습에 의존하지 않고도 기존 베이스라인 대비 유의미한 성능 향상을 달성하였습니다.

원저자: Rutav Shah, Rajat Kumar Jenamani, Xiaohan Zhang, Lingfeng Sun, Roberto Martín-Martín, Yuke Zhu, Deva Ramanan, Karl Schmeckpeper

게시일 2026-06-16

📖 4 분 읽기☕ 가벼운 읽기

CC BY 4.0

원저자: Rutav Shah, Rajat Kumar Jenamani, Xiaohan Zhang, Lingfeng Sun, Roberto Martín-Martín, Yuke Zhu, Deva Ramanan, Karl Schmeckpeper

원본 논문은 CC BY 4.0 (http://creativecommons.org/licenses/by/4.0/) 라이선스로 제공됩니다. ✨ 이것은 아래 논문에 대한 AI 생성 설명입니다. 저자가 작성하거나 승인한 것이 아닙니다. 기술적 정확성을 위해서는 원본 논문을 참조하세요. 전체 면책 조항 읽기

복잡한 요리를 하려고 하는데, 양념을 집으려고 고개를 돌릴 때마다 방금 무엇을 하려 했는지 즉시 까먹는 상황을 상상해 보세요. 가스레인지를 켰지만, 팬을 가져오기 위해 자리를 비운 사이 끄는 것을 잊어버리거나, 심지어 소금을 이미 넣었다는 사실조차 잊어버릴 수도 있습니다. 이것이 바로 현재 많은 로봇이 직면한 문제입니다. 그들은 눈앞에 있는 것을 보는 데는 매우 똑똑하지만, 몇 초 또는 몇 분 전에 일어난 일을 기억하는 '단기 기억력'은 거의 없습니다.

이 논문은 이 문제를 해결하기 위해 새로운 로봇의 뇌인 PRISM과 새로운 테스트인 REMEMBENCH를 소개합니다. 작동 방식은 다음과 같이 쉽게 설명할 수 있습니다.

문제점: "금붕어 기억력"을 가진 로봇들

대부분의 현대적 로봇은 인간이 하는 행동을 관찰하며 배웁니다(마치 학생이 선생님을 따라 하는 것과 같습니다). 하지만 이들은 보통 현재의 비디오 프레임 하나만을 봅니다. 만약 작업이 "가스레인지를 켰으니, 끄기 전에 2분을 기다려야 한다"와 같이 30초 전의 일을 기억해야 하는 것이라면, 로봇은 과거의 행동을 기억하지 못하기 때문에 실패하게 됩니다.

이를 해결하기 위해 단순히 로봇에게 더 긴 비디오 기록을 보여주도록 만들면, 두 가지 나쁜 결과가 발생합니다.

노이즈 문제 (The Noise Problem): 로봇은 너무 많은 정보에 압도당합니다. 로봇은 배경에 있는 무관한 물체 같은 쓸모없는 세부 사항을 자신의 결정과 연결하기 시작하여 실수를 유발할 수 있습니다.
과도한 연산 문제 (The Heavy Lifting Problem): 긴 비디오 기록을 처리하는 것은 엄청난 컴퓨터 성능과 메모리를 요구하며, 이는 로봇을 느리고 비싸게 만듭니다.

해결책: PRISM (스마트한 사서)

저자들은 단순한 수집가가 아닌 스마트한 사서처럼 행동하는 로봇 정책인 PRISM을 구축했습니다. 모든 페이지를 머릿속에 다 담아두는 대신, PRISM은 두 가지 영리한 기술을 사용합니다.

"게이트 어텐션" 필터 (경비원):
로봇의 기억을 붐비는 방이라고 상상해 보세요. PRISM에는 누가 들어오고 누가 나갈지를 결정하는 경비원이 있습니다. 만약 로봇이 "빨간 사과를 집었다"는 것을 기억하고 있지만, 현재 작업이 파란 그릇에 관한 것이라면, 경비원은 빨간 사과의 기억이 로봇을 산만하게 하지 못하도록 차단합니다. 이를 통해 로사적인 과거 사건과 현재 행동 사이의 어처구니없는 연결을 막아줍니다.
"계층적" 요약기 (뉴스 앵커):
긴 이야기의 모든 단어를 하나하나 읽는 대신, PRISM은 먼저 작은 덩어리들을 읽고 각 덩어리에 대한 짧은 요약을 작성합니다. 그런 다음, 전체적인 줄거리를 이해하기 위해 그 요약본들을 읽습니다. 이는 뉴스 앵커가 모든 원본 경찰 보고서를 읽는 대신 그날의 사건을 요약하는 것과 같습니다. 이 방식 덕분에 로봇은 훨씬 빨라지고 컴퓨터 메모리도 적게 사용하며, 최대 2분의 기록을 기억할 수 있습니다.

테스트: REMEMBENCH (로봇 기억력 시험)

저자들은 로봇이 실제로 기억력을 가지고 있는지 증명하기 위해 REMEMBENCH를 만들었습니다. 이것은 로봇의 기억력을 측정하는 표준화된 "운전면허 시험"이라고 생각하면 됩니다. 단순히 팔을 움직일 수 있는지를 테스트하는 것이 아니라, 네 가지 특정 유형의 단기 기억력을 테스트합니다.

공간 기억 (Spatial Memory): "내가 과일을 어디에 두었지?" (로봇은 현재 보이지 않는 물체를 찾아내야 합니다).
미래 계획 기억 (Prospective Memory): "2분 후에 가스레인지를 꺼야 한다." (로봇은 과거의 트리거를 바탕으로 미래의 행동을 기억해야 합니다).
대상 연관 기억 (Object-Associative Memory): "이 사과를 씻었으니, 반드시 이 특정 그릇에 다시 담아야 한다." (로봇은 어떤 물체가 어떤 용기와 짝인지 기억합니다).
대상 집합 기억 (Object-Set Memory): "빵스틱을 모두 옮겨야 하지만, 옮기는 동안 개수를 세어야 한다." (로봇은 물체 그룹의 합계를 계속 추적합니다).

결과: 명확한 승자

저자들이 PRISM을 다른 로봇들(표준적인 "긴 기억"을 가진 로봇이나 다른 유형의 AI 뇌를 가진 로봇 포함)과 비교 테스트했을 때의 결과입니다.

기억력 테스트에서: PRISM은 압도적인 승자였으며, 그다음으로 뛰어난 로봇보다 상당한 차이(5%~12%)로 앞섰습니다.
표준 작업에서: 기억력을 명시적으로 테스트하지 않는 작업에서도 PRISM은 더 나은 성과를 보였습니다. 왜냐하면 기억을 갖는 것은 로봇이 맥락을 이해하는 데 도움을 주기 때문입니다. 예를 들어, 컵을 "방금 집어 들었는지" 아니면 "놓으려는 중인지" 아는 것은 혼란을 피하는 데 도움이 됩니다.
실제 환경: 실제 물리적 로봇에 적용했을 때, 기억이 없는 로봇은 완전히 실패(성공률 0%)한 반면, PRISM은 30%의 성공률을 보였습니다.

핵심 요약

이 논문은 "경비원"을 통해 노이즈를 걸러내고 "요약기"를 통해 공간을 절약함으로써, 로봇이 몇 분 전의 일을 마침내 기억할 수 있다고 주장합니다. 이를 통해 로봇은 단순히 현재 순간에 반응하는 것을 넘어, 일련의 단계가 필요한 길고 복잡한 가사 노동을 수행할 수 있습니다. 저자들은 더 나은 기억 능력을 갖춘 로봇을 만들 수 있도록 로봇의 뇌(PRISM)와 테스트(REMEMBENCH)를 제공합니다.

기술 요약: 장기 과업을 위한 시각-운동 정책의 단기 기억 확장

1. 문제 정의

로봇 과업은 부분 관측성(partial observability)을 처리하기 위해 단기 기억(수 초에서 수 분에 이르는 범위)을 필요로 하는 경우가 많습니다. 예를 들어, 더 이상 보이지 않는 물체를 회수하거나, 가전제품을 끄는 것과 같이 동작의 타이밍을 맞추는 작업 등이 이에 해당합니다. 그러나 모방 학습을 통해 훈련된 대부분의 현재 트랜스포머 기반 정책들은 과거의 경험을 활용하지 못한 채 즉각적인 감각 입력에만 의존하는 **근시안적(myopic)**인 특성을 보입니다.

이러한 문제를 해결하기 위해 표준 트랜스포머 기반 정책의 컨텍스트 윈도우를 확장하여 긴 맥락을 포함하도록 만드는 데에는 두 가지 핵심적인 도전 과제가 존재합니다:

가짜 상관관계(Spurious Correlations): 단순히 긴 이력을 참조하게 되면 무관한 정보(예: 과거의 방해 요소 위치)가 도입되어, 모델이 잘못된 상관관계를 학습하게 함으로써 성능을 저하시킵니다.
계산 효율성(Computational Scalability): 표준 어텐션 메커니즘의 계산 및 메모리 점유율은 컨텍스트 길이에 따라 선형 또는 이차적으로 증가합니다. 고차원 시각 입력을 긴 시간적 시퀀스(예: 몇 분 단위)에 대해 처리하는 것은 비용이 매우 많이 듭니다.

2. 방법론: PRISM

저자들은 단기 기억을 효과적으로 활용하기 위해 설계된 트랜스포머 기반 아키텍처인 PRISM(Policy Architecture with Short-Term Memory)을 제안합니다. 이 모델은 두 가지 핵심 혁신을 특징으로 합니다:

게이트형 어텐션 (정보 필터링): 가짜 상관관계를 완화하기 위해, PRISM은 **사후 어텐션 게이팅 메커리즘(post-attention gating mechanism)**을 도입합니다. 이 구성 요소는 현재 입력 특징에 조건화된 학습 가능한 게이트를 사용하여 각 어텐션 블록의 출력을 선택적으로 필터링합니다. 이를 통해 정책은 이력으로부터 오는 무관한 세부 정보를 적응적으로 억제하여 동작 예측의 노이즈를 줄일 수 있습니다.
계층적 구조 (효율성): 확장성 문제를 해결하기 위해 PRISM은 계층적 설계를 채택합니다.
- 로컬 인코딩(Local Encoding): 로컬 메모리 인코더가 감각 입력을 처리하고, Perceiver 스타일의 리샘플러를 사용하여 이를 압축된 토큰으로 압축합니다.
- 글로벌 어텐션(Global Attention): 글로벌 메모리 인코더가 이러한 압축된 토큰들에 대해 시간 축을 가로질러 어텐션을 계산합니다.
- 융합(Fusion): 생성된 글로벌 정보는 잔차 연결(residual connections)을 통해 원래의 로컬 토큰과 결합되어 다시 브로드캐스트됩니다.
- 이점: 이 방식은 계산 복잡도를 $O((n \cdot k_{tot})^2)$ 에서 약 $O(n \cdot k_{tot} \cdot m + (n \cdot k_{tot}/m)^2)$ 로 줄여줍니다 (여기서 $m$ 은 압축 계수). 또한 테스트 시점에 압축된 요약 토큰만을 캐싱함으로써 메모리 점유율을 크게 줄입니다.

PRVM은 추가적인 손실 함수나 외부 메모리 모듈 없이, 폐루프(closed-loop) 환경에서 과거의 감각 입력과 동작의 시퀀스를 미래의 동작으로 매핑하는 표준 행동 복제(behavior cloning)(모방 학습)를 사용하여 훈련됩니다.

3. 주요 기여

PRISM 아키텍처: 게이트형 어텐션과 계층적 요약을 결합하여, 계산 효율성과 노이즈에 대한 강건성을 유지하면서 단기 기억을 최대 2분(2~3 Hz 기준)까지 확장할 수 있는 새로운 정책 아키텍처를 제시합니다.
REMEMBENCH: 인지 과학에서 유래한 네 가지 범주의 단기 기억을 포괄하는 8가지 다양한 가전 조작 과업으로 구성된 새로운 벤치마크입니다:
- 공간적(Spatial): 물체의 위치를 기억함 (예: 숨겨진 과일 회수).
- 미래 지향적(Prospective): 지연 시간 동안 의도를 유지함 (예: 특정 시간 동안 고기 요리하기).
- 객체 연상(Object-Associative): 객체와 위치 간의 연상을 기억함 (예: 씻은 물건을 원래의 특정 접시로 되돌려 놓기).
- 객체 집합(Object-Set): 여러 물체의 집합을 유지하고 업데이트함 (예: 전자레인지로 옮긴 브레드스틱 개수 세기).
- REMEMBENCH는 단일 측면이나 간접적인 기억 테스트에 집중하는 기존 벤치마크와 달리, 즉각적인 감각 입력만으로는 성공할 수 없도록 설계되어 명시적으로 기억을 요구합니다.

4. 실험 결과

본 논문은 REMEMBENCH 및 표준 벤치마크에서 PRISM을 순환형 아키텍처(LSTM, Mamba), 세그먼트 수준 순환 트랜스포머(Transformer-XL), 그리고 기타 트랜스포머 변형(Linear Attention, Scene Memory Transformer)과 비교 평가합니다.

REMEMBENCH 성능: PRISM은 평균 성공률 **41%**를 달기하며, 가장 강력한 베이스라인들(SMT 36%, PTP 28%, Mamba 25%)보다 절대 수치로 5%~12% 높은 성능을 보였습니다.
- 확장성: 성능은 메모리 크기에 따라 확장됩니다. 윈도우를 1에서 512 타임스텝으로 늘리면 성공률이 16%에서 42%로 향상되며 포화 상태에 도달하지 않습니다.
- 절제 실험(Ablation): 게이팅 메커니즘을 제거하면 성능이 16포인트 하락하며, 이는 무관한 이력을 필터링하는 데 있어 게이팅의 필요성을 입증합니다.
표준 벤치마크 (ROBOCASA & LIBERO):
- ROBOCASA에서 메모리를 포함한 PRISM(n=256)은 **43%**의 성공률을 기록하며, 메모리가 없는 변형 모델보다 11%, 미세 조정된 GR00T-N1-2B 베이스라인보다 **11%**의 절대적 개선을 보였습니다.
- LIBERO에서 PRISM은 평균 **89%**의 성공률을 달성하며, 미세 조정된 OpenVLA 베이스라인보다 15%, 메모리가 없는 변형 모델보다 12% 높은 성능을 보였습니다.
- 저자들은 메모리를 추가하는 것이 시각적으로 동일한 상태(예: "방금 잡은 상태"와 "놓으려는 상태"의 구분)에서의 동작 모호성을 해결하여, 메모리 테스트용으로 설계되지 않은 과업에서도 다중 양상성(multimodality)을 줄여준다고 언급했습니다.
실세계 평가: "세척 후 용기에 되돌려 놓기" 과업의 실세계 적용에서, PRISM은 **30%**의 성공률을 기록했습니다. 이는 메모리가 없는 트랜스포머 베이스라인의 15%, 메모리가 없는 정책의 **0%**와 대조됩니다.
효율성: 컨텍스트 길이가 512 스텝일 때, PRISM은 전체 글로벌 어텐션 및 하이브리드 슬라이딩 윈도우 방식에 비해 피크 GPU 메모리 사용량을 80% 줄이고 실행 시간을 40% 단축했습니다.

5. 의의 및 주장

본 논문은 PRISM과 REMEMBENCH가 장기 과업을 수행할 수 있는 단기 기억 증강 시각-운동 정책을 개발하고 평가하기 위한 기초를 마련한다고 주장합니다.

일반화: REMEMBENCH는 네 가지 뚜렷한 기억 범위를 다룸으로써, 특정 과업에 국한된 솔루션이 아닌 일반적인 메모리 메커니즘의 발전을 독려합니다.
확장성: PRISM은 단순한 컨텍스트 윈도우 확장에 따른 과도한 비용 없이도 단기 기억을 수 분의 상호작용까지 확장할 수 있음을 보여줍니다.
강건성: 게이트형 어텐션 메커니즘은 노이즈가 많은 장기 데이터에서 모델이 가짜 상관관계를 학습하는 것을 방지하는 데 필수적임을 입증했습니다.

저자들은 REMEMBENCH에서의 현재 성공률(약 30~40%)이 장기 과업이 여전히 어렵다는 것을 나타내지만, 다양한 아키텍처와 벤치마크 전반에서 나타나는 일관된 메모리 증강 효과가 이 접근법의 타당성을 검증한다고 결론짓습니다. 그들은 지속적인 장기 기억의 통합 및 메모리 모듈을 위한 인터넷 규모의 사전 훈련 활용 등을 향후 연구 방향으로 제시하지만, 현재의 작업은 모방 학습 프레임워크 내에서의 단기 기억 메커로니즘에 엄격히 집중하고 있음을 강조합니다.