기술 요약: TaskSense

문제 정의

시각적 제어를 위한 월드 모델(World models)은 일반적으로 압축된 잠재 상태(latent states)를 학습하기 위해 전체 관측값(observations)을 재구성하는 방식에 의존한다. 이러한 방식은 정적인 배경을 가진 표준화된 벤치마크에서는 효과적이지만, 잠재 표현이 제어 작업과 무관한 전체 시각적 입력 정보까지 보존하도록 암묵적으로 유도한다. 작업 관련 콘텐츠가 관측값의 아주 작은 부분만을 차지하거나 배경의 노이즈 및 방해 요소(distractors)가 존재하는 환경에서, 이러한 재구성 목적 함수는 작업과 무관한 시각적 콘텐츠에 귀중한 표현 용량을 소모하게 만든다. 시각적 재구성과 제어 목적 사이의 이러한 불일치는 잠재 역학(latent dynamics)을 편향시켜, 제어 관련 특징을 위한 학습 신호를 희석시키고 시각적 방해 요소가 있을 때 다운스트림 성능을 심각하게 저하시킨다. 반대로, 오직 작업 감독(예: 보상 예측)에만 의존하는 것은 풍부한 시각적 표현을 학습하기에는 신호가 너무 약한 경우가 많다. 재구성을 완전히 제거하는 기존 솔루션들은 상상 기반 계획(imagination-based planning)에 필요한 생성 능력을 희생하며, 특징 추출 후에 어텐션이나 마스킹을 적용하는 방법들은 애초에 어떤 정보가 잠재 역학 모델로 들어올지를 제어하지 못한다는 한계가 있다.

방법론

저자들은 TaskSense를 제안한다. 이는 잠재 인코딩 전에 작업 관련성을 강제하는 작업 중심(task-centric) 월드 모델링 프레임-워크이다. TaskSense는 DreamerV3 아키텍처에 이전 잠재 믿음 상태(previous latent belief state)에 의해 조건화된 미분 가능한 확률적 공간 어텐션(differentiable stochastic spatial attention) 메커니즘을 추가한다.

핵심 구성 요소

1. 잠재 조건부 확률적 공간 어텐션 (Latent-Conditioned Stochastic Spatial Attention):

   • 모든 픽셀에 대해 어텐션을 예측하는 대신, 모델은 관측 범위를 덮는 조밀한 격자(lattice) $\Lambda$에 대해 어텐션을 정의한다.
   • 경량 MLP는 이전 잠재 상태($s_{t-1}$, 구체적으로 은닉 상태 $h_{t-1}$와 확률적 잠재 변수 $z_{t-1}$)를 기반으로 어텐션 로짓(logits)을 예측하며, 이를 통해 모델은 즉각적인 시각적 외형이 아니라 진화하는 작업 이해도를 바탕으로 어디에 집중할지 결정할 수 있게 한다.
   • 어텐션 필드는 Binary Concrete (Relaxed Bernoulli) 분포를 사용하여 독립적인 확률 변수들의 집합으로 모델링된다. 이를 통해 모델은 여러 영역을 동시에 주목하거나 무시할 수 있으며, 소프트맥스(softmax) 어텐션의 경쟁적 제약을 피할 수 있다.
   • 학습 중에는 미분 가능한 마스크가 샘플링되며, 추론 중에는 결정론적 평균값이 사용된다.

2. 어텐션 인식 인코딩 (Attention-Aware Encoding):

   • 샘플링된 어텐션 마스크는 업샘플링된 후 요소별 곱셈(element-wise multiplication)을 통해 관측값을 변조하여, 낮은 어텐션을 받은 영역을 감쇄시킨다.
   • 낮은 강도의 픽셀과 억제된 높은 강도의 픽셀 사이의 모호함을 해결하기 위해, 업샘플링된 어텐션 맵이 인코더의 추가 입력 채널로 결합된다.

3. 보조 역역학 목적 함수 (Auxiliary Inverse-Dynamics Objective):

   • 희소성 정규화(sparsity regularizer)가 어텐션을 모든 관측을 무시하는 사소한 해(trivial solution)로 수렴시키는 것을 방지하기 위해 보조 목적 함수를 도입한다.
   • 모델은 인코딩된 특징의 시간 차이($f_t - f_{t-1}$)를 사용하여 두 연속적인 타임 스텝 사이에 실행된 행동 $a_{t-1}$을 예측한다.
   • 이 목적 함수는 어텐션 메커니즘을 제어 관련 영역에 고정시켜, 에이전트의 행동을 예측하는 데 유용한 시각적 특징을 유지하도록 유도한다.

4. 어텐션 조건부 재구성 (Attention-Conditioned Reconstruction):

   • 전체 관측값을 재구성하는 대신, 오직 주목된 영역($o^{att}_t$)만을 재구성한다.
   • 디코더는 샘플링된 어텐션 맵(stop-gradient 적용)에 조건화되어, 의도적으로 필터링된 정보에 대해 잠재 상태가 페널티를 받지 않도록 함으로써 어떤 영역이 유지되었는지 인지하게 한다.

학습 목적 함수

전체 손실 함수는 표준 DreamerV3의 잠재 역학 손실($L_{dyn}$)과 새로운 구성 요소들을 결합한다:
$$L_{WM} = L_{dyn} + L_{rec} + \lambda_{sparse}L_{sparse} + \lambda_{inv}L_{inv}$$
여기서 $L_{rec}$는 필터링된 관측값에 대한 재구성 손실이며, $L_{sparse}$는 압축된 어텐션 맵을 권장하고, $L_{inv}$는 역역학 손실이다.

주요 기여

1. TaskSense 프레임워크: 잠재 인코딩 전 단계에서 잠재 조건부 확률적 공간 어텐션을 통해 작업 관련성을 강제하는 재구성 기반 월드 모델링 프레임워크를 제시한다.
2. 자기 지도 학습 메커니즘: 역역학 감독과 필터링된 관측값에 대한 재구성을 결합하여 확률적 어텐션을 제어 관련 영역에 고정하는 새로운 접근 방식을 제공한다.
3. 실증적 검증: TaskSense가 표준 벤치마크에서 경쟁력 있는 성능을 유지하면서도 시각적 방해 요소가 있는 환경에서 베이스라인보다 크게 향상된 성능을 보임을 입증한다.

실험 결과

저자들은 **Deepism Control Suite (DMC)**와 동적 자연 이미지 배경을 특징으로 하는 **Distracting Control Suite (DCS)**에서 TaskSense를 평가하였다.

• 표준 성능: 정적인 배경을 가진 표준 DMC에서 TaskSense는 대부분의 작업에서 DreamerV3 베이스라인과 대등한 성능을 보여주었으며, 이는 관측값의 일부만을 재구성하는 것이 효과적인 잠재 계획에 충분함을 입증한다.
• 방해 요소에 대한 강건성: Distracting Control Suite에서 TaskSense는 일관되게 DreamerV3를 능가한다. 예를 들어, Walker Run 작업에서 TaskSense는 평균 리턴 531.2를 기록하여 DreamerV3의 286.2를 상회했다. 이는 시각적 무관 콘텐츠를 필터링하는 것이 강건성을 향상시킨다는 것을 확인시켜 준다.
• 정성적 분석: Cheetah Run 작업에서 학습된 어텐션 맵을 시각화하면, TaskSense가 몸통과 방해되는 배경을 억제하고 제어 관련 영역(예: 치타의 다리)에 일관되게 집중하는 것을 볼 수 있다.
• 절제 연구 (Ablation Studies):
  • 역역학 목적 함수를 제거하면 어텐션이 분산되어 제어 관련 영역을 국지화하지 못하고, 이로 인해 상당한 성능 저하가 발생한다.
  • 확률적 어텐션을 결정론적 어텐션으로 교체하면 성능이 크게 하락하며, 이는 훈련 과정에서 모델이 다양한 부분적 관측값에 노출되는 것의 중요성을 강조한다.
  • 어텐션 조건부 재구성을 제거하면 다소 작지만 눈에 띄는 성능 저하가 발생하며, 이는 재구성 대상과 필터링된 입력을 일치시키는 것의 필요성을 나타낸다.

의의 및 주장

본 논문은 현재의 월드 모델이 가진 근본적인 불일치, 즉 제어에 중요한 것보다 모든 시각적 정보를 보존하려는 경향을 TaskSense가 해결한다고 주장한다. Task-relevance를 사후적인 표현 정제(post-hoc representation refinement)가 아닌 관측 레벨의 필터링을 통해 강제함으로써, TaskSense는 재구성 기반 월드 모델이 제어 관련 정보에 효율적으로 표현 용량을 할당할 수 있게 한다.

저자들은 이 접근 방식이 재구성의 생성적 이점(상상 기반 계획에 필수적)을 유지하면서도, 일반적인 모델들을 저하시키는 시각적 방해 요소에 대한 강건성을 확보할 수 있게 한다고 단언한다. 이 연구는 작업 관련성을 관측 인터페이스 단계에서부터 강제하는 것이 시각적 월드 모델과 표현 학습을 개선하는 유망한 방향임을 시사한다.

논문은 한 가지 한계를 인정한다: 시각적 특징이 즉각적인 행동을 예측하는 것(역역학에 의해 포착됨)이 장기적인 작업 성공에 유효한 것과 동일하다는 가정을 전제로 한다는 점이다. 두 요소가 서로 어긋나는 시나리오에서는 현재의 감독 신호가 최적이 아닌 특징에 주의를 집중시킬 수 있으며, 이는 향후 연구 과제로 식별되었다.