기술 요약: 시각-운동 정책 학습을 위한 순서형 액션 토큰 (Ordered Action Tokens)

문제 정의

액션 토큰화는 현대의 시각-운동 정책(visuomotor policies)에서 연속적인 로봇 제어와 시퀀스 모델링 사이를 잇는 핵심적인 인터페이스 역할을 합니다. 그러나 기존 방식들은 효과적인 정책 학습에 필수적인 세 가지 속성을 동시에 만족시키지 못합니다:

1. 높은 압축률 (High Compression): 학습의 난이도와 추론 지연 시간을 최소화하기 위해 토큰 시퀀스는 조밀해야 합니다.
2. 완전한 디코딩 가능성 (Total Decodability): 정책이 생성한 임의의 토큰 시퀀스는 반드시 유효하고 실행 가능한 연속적 액션으로 신뢰성 있게 디코딩되어야 합니다.
3. 순서형 토큰 구조 (Ordered Token Structure): 토کن 시퀀스는 제어 관련 정보를 조직화하여, 초기 토큰은 거시적이고 영향력이 큰 움직임을 포착하고, 후기 토큰은 잔여 세부 사항을 정교화해야 합니다. 이러한 "조대-정밀(coarse-to-fine)" 구조는 자기회귀 모델의 차기 토큰 예측(next-token prediction) 유도 편향과 일치하며, 시각-언어 모델(VLM)에 의미 있는 감독(supervision)을 제공합니다.

현재의 방법들은 이러한 요구사항 중 일부만을 만족할 뿐, 세 가지를 동시에 충족하지는 못합니다. 차원별 빈닝(per-dimension binning)은 완전한 디코딩이 가능하지만 지나치게 긴 시퀀스를 생성합니다. 주파수 영역 토커나이저(예: FAST)는 순서 구조를 제공하지만, 가변 길이 확장 문제로 인해 완전한 디코딩 능력이 부족합니다. 학습된 잠재 토커나이저(예: QueST, ACodec)는 압축률은 높지만, 전역 제어 정보를 초기에 배치하는 명시적인 순서 구조가 결여된 경우가 많아 정책 타겟으로서의 효과가 제한적입니다.

방법론: 순서형 액션 토큰화 (Ordered Action Tokenization, OAT)

저자들은 세 가지 요구사항을 모두 만족하도록 설계된 학습된 토커나이저인 **OAT (Ordered Action Tokenization)**를 제안합니다. OAT는 연속적인 액션 청크를 조밀하고, 완전히 디코딩 가능하며, 순서가 있는 토큰 시퀀스로 이산화합니다.

아키텍처

OAT는 이산적 병목(discrete bottleneck)을 갖춘 인코더-디코더 아키텍처를 채택합니다:

• 레지스터 병목 (Register Bottleneck): 인코더는 연속적인 액션 청크의 모든 위치에 교차 주의(cross-attend)하기 위해 학습 가능한 "레지스터" 토큰(ViT 레지스터에서 영감을 받음)을 사용합니다. 이 레지스터들은 시간적 정보를 고정된 길이의 잠재 상태(latent states) 시퀀스로 집계합니다.
• 유한 스칼라 양자화 (Finite Scalar Quantization, FSQ): 잠재 레지스터 상태는 FSQ를 통해 고정된 어휘집(vocabulary)으로 이산화되며, 이를 통해 모든 어휘집 인덱스가 유효한 연속 액션 청크로 매핑됨을 보장합니다 (완전한 디코딩 가능성).
• 디코더 (Decoder): 교차 주의 트랜스포머 디코더는 양자화된 레지스터 상태로부터 연속 액션 청크를 재구성합니다.

학습: 순서형 접두사 재구성 (Ordered Prefix Reconstruction)

원하는 조대-정밀 순서를 유도하기 위해, OAT는 학습 과정에서 **중첩된 접두사 재구성 (Nested Prefix Reconstruction)**을 활용합니다:

• 중첩 드롭아웃 (Nested Dropout): 학습 목적 함수는 여러 예산(budget) 집합(예: $\{1, 2, 4, \dots, H_l\}$)에서 재구성 예산 $K$를 샘생합니다. 디코더는 나머지 접미사(suffix)를 마스킹 처리한 상태에서 오직 처음 $K$개의 토큰만을 사용하여 전체 액션 청크를 재구성하도록 학습됩니다.
• 정보 할당 (Information Allocation): 이는 초기 토큰이 그럴듯한 액션을 수행하는 데 필요한 거시적이고 전역적으로 중요한 제어 정보를 운반하도록 강제하며, 후기 토큰은 잔여 세부 사항을 정교화하게 만듭니다. 이는 어떤 접두사라도 실행 가능한 액션으로 디코딩될 수 있는 "점진적" 표현을 생성합니다.
• 레지스터 흐름 제약 (Register Flow Constraints): 저자들은 다양한 생성 스케줄을 지원하기 위해 레지스터를 위한 특정 셀프 어텐션 마스크를 도입했습니다. 두 가지 변형을 정의합니다:
  • $OAT_{sing}$ (토큰 단위): 표준 인과적 주의(causal attention)를 사용하여 모든 토큰에 엄격한 의존 순서를 부여합니다.
  • $OAT_{pow2}$ (2의 거듭제곱 단위): 블록-인과적 주의(block-causal attention)를 사용하여, 2의 거듭제곱 그룹 내의 레지스터들은 서로를 참조할 수 없고 오직 이전 그룹들만 참조할 수 있습니다. 이는 확장 가능한 블록 자기회귀 디코딩을 가능하게 합니다.

정책 인터페이스

OAT는 두 가지 주요 정책 패러다임에서 검증되었습니다:

1. 자기회귀 (Autoregressive, AR) 정책: 정책이 OAT 토큰을 순차적(또는 블록 단위)으로 생성하면, 이를 디토크나이징하여 액션으로 변환합니다. 이 순서 구조는 가변적인 추론 예산을 지원하여, 시스템이 추론 비용과 액션 충실도 사이에서 균형을 맞출 수 있게 합니다.
2. 토큰 공동 학습 (Token Co-Training, TC) 정책: VLM은 OAT 토큰을 예측하도록 학습되는 반면(감독), 별도의 플로우 매칭(flow-matching) 전문가가 분리된 VLM 컨텍스트로부터 연속 액션을 생성합니다. 여기서 첫 번째 OAT 토큰은 전체 액션 청크의 "계획과 같은" 요약본 역할을 하여, VLM의 프리필(prefill) 표현에 강력한 전역적 목적을 제공합니다.

확장 가능한 생성

본 논문은 토큰 단위 생성, 원샷 병렬 디코딩, 중간 블록 스케줄을 통합하는 블록 단위 자기회귀 (Block-wise Autoregressive, BAR) 프레임워크를 소개합니다. 구체적으로, $OAT_{pow2}$ 변형을 2의 거듭제곱 생성 스케줄과 결합하면 순차적 정책 호출 깊이를 $O(H_l)$에서 $O(\log H_l)$로 줄여, 성능 저하 없이 추론 효율성을 크게 향고할 수 있습니다.

주요 기여

1. 요구사항의 공식화: 본 논문은 높은 압축률, 완전한 디코딩 가능성, 그리고 순서형 토큰 구조를 정책 인터페이스로서의 액션 토큰화에 대한 세 가지 핵심 요구사항으로 식별하고 공식화했습니다.
2. OAT 알고리즘: 레지스터 기반 압축, FSQ, 그리고 중첩 접두사 학습을 통해 이러한 요구사항을 충족하는 학습된 토커나이저를 제안했습니다.
3. 확장 가능한 디코딩 프레임워크: 다양한 생성 스케줄을 통합하는 BAR 프레임워크를 도입하여, 2의 거듭제곱 그룹화가 어떻게 추론 지연 시간을 획기적으로 줄일 수 있는지 입증했습니다.
4. 포괄적인 검증: 세 가지 정책 백본(경량 트랜스포머, PaliGemma2, Qwen3VL)과 AR 및 TC 패러다임을 아우르는 60개 이상의 작업(시뮬레이션 및 실제 환경)에 대해 광범위한 실험을 수행했습니다.

실험 결과

• 율-왜곡 (Rate-Distortion): OAT는 재구성 오차와 토큰 예산 사이의 매끄러운 곡선을 추적하며, 다양한 길이의 접두사가 점진적으로 더 정확한 액션으로 디코딩됨을 보여줍니다.
• 경량 정책: 고정된 트랜스포머 백본을 사용하는 폐루프 제어에서, OAT는 베이스라인(Bin, FAST, QueST, ACodec)을 일관되게 능가합니다. 성능은 토큰 예산에 따라 확장되며, $OAT_{sing}^{8}$은 LIBERO, RoboMimic, MetaWorld 및 실제 작업에서 가장 높은 성공률을 달s합니다.
• 자기회귀 VLM: OAT는 대규모 VLM(PaliGemma2, Qwen3VL)의 성공률을 향상시킵니다. 2의 거듭제곱 변형($OAT_{pow2}$)은 토큰 단위 생성과 대등한 성능을 보이면서도 훨씬 적은 횟수의 정책 호출(예: 전체 생성을 위한 16회 대신 5회)을 필요로 합니다.
• 토큰 공동 학습 (Token Co-Training): 추론 시 토큰을 디코딩하지 않더라도, OAT 감독은 우수한 VLM 표현을 이끌어냅니다. OAT는 TC 설정에서 여러 벤치마크에서 QueST와 ACodec을 능가하며 가장 높은 평균 성공률을 달성하여, "계획과 같은" 감독 가설을 검증했습니다.
• 어블레이션 (Ablations): 순서형 접두사 학습(중첩 드롭아웃)을 제거하면 성능이 크게 저하되어, 조대-정밀 구조가 매우 중요하다는 것을 확인했습니다. 또한, 레지스터 어텐션 구조와 생성 스케줄(예: $OAT_{sing}$을 2의 거듭제곱 블록과 함께 사용)이 불일치할 경우 성능이 감소하여, 토커나이저 설계와 추론 스케줄 간의 정렬이 필요함을 보여주었습니다.

의의 및 주장

본 논문은 OAT가 시각-운동 정책 학습을 위한 견고하고 유연한 인터페이스를 제공한다고 주장합니다. 압축, 디코딩 가능성, 순서라는 삼박자를 만족함으로써 OAT는 다음을 가능하게 합니다:

• 유연한 추론: 정책이 지연 시간과 충실도 사이에서 균형을 맞추기 위해 적응적으로 토큰 예산을 선택할 수 있습니다.
• 학습 개선: 순서 구조는 자기회귀 모델의 유도 편향과 일치하며, 공동 학습 설정에서 VLM에 의미 있는 전역적 감독을 제공합니다.
• 확장성: BAR와의 통합을 통해, 선형적인 토큰 단위 생성의 깊이 페널티 없이 대규모 모델에서 효율적인 추론을 가능하게 합니다.

저자들은 본 연구가 특정 벤치마크와 백본에 국한되어 있으나, OAT의 설계 원칙이 액션 토큰화 문제에 대한 일반화 가능한 솔루션을 제공하며, 더 효율적이고 유능한 체화된 AI(embodied AI) 시스템으로 나아가는 길을 열어준다고 결론짓습니다. 또한, 향후 연구에서 과업의 복잡성이나 불확실성에 따라 정책이 필요한 토큰 예산을 동적으로 결정하는 적응형 계산(adaptive computation)을 탐구할 수 있음을 언급합니다.