기술 요약: 외생적 컨텍스트 MDP 학습을 위한 Minimax PAC 경계

1. 문제 설정

본 논문은 외생적 i.i.d. 컨텍스트를 가진 테이블형 할인 마르코프 결정 과정(tabular discounted MDPs) 학습의 샘플 복잡도를 조사한다. 모델은 튜플 $\mathcal{M} = (\mathcal{X}, \mathcal{Z}, \mathcal{A}, P, r, \mu, \gamma)$로 정의된다:

• $\mathcal{X}$는 유한한 제어 상태 공간이다.
• $\mathcal{A}$는 유한한 행동 공간이다.
• $\mathcal{Z}$는 유한한 컨텍스트 공간이다.
• $\gamma \in [0, 1)$는 할인 인자이다.
• $P: \mathcal{X} \times \mathcal{A} \times \mathcal{Z} \to \Delta(\mathcal{X})$는 전이 커널이다.
• $r: \mathcal{X} \times \mathcal{A} \times \mathcal{Z} \to [0, 1]$은 보상 함수이다.
• $\mu$는 알려지지 않은 컨텍스트 분포이다.

각 타임스텝 $t$에서, 컨텍스트 $Z_t$는 $\mu$로부터 독립적으로 추출되어 행동 선택 전에 에이전트에게 공개된다. 에이전트는 $A_t$를 선택하고, 보상 $r(X_t, A_t, Z_t)$를 받으며, $X_{t+1} \sim P(\cdot | X_t, A_t, Z_t)$로 전이된다. 결정적으로, 다음 컨텍스트 $Z_{t+1}$은 에이전트의 행동 및 현재 상태와 무관하며, 오직 $\mu$에만 의존한다.

학습 목표는 정책 평가(Policy Evaluation, PE), 최적 가치 추정(Best-Value Estimation, BVE), 그리고 **최적 정책 추출(Best-Policy Extraction, BPE)**이다. 샘플 복잡도는 오프라인 학습 단계에서 수행된 오라클 호출 횟수와 의사결정 시점(쿼리) 단계에서의 호출 횟수를 나타내는 쌍 $(n_{\text{learn}}, m_{\text{query}})$로 측정된다. 본 논문은 두 가지 정칙(regime)을 고려한다:

1. 알려진 역학(Known Dynamics): $P$와 $r$은 알려져 있고, 오직 $\mu$만이 알려져 있지 않다.
2. 완전히 알려지지 않은 경우(Fully Unknown): $\mu$와 $P$가 모두 알려져 있지 않다.

2. 방법론

2.1 핵심 통찰: 평균 가치 함수

본 연구의 핵심적인 방법론적 기여는 문제를 확장된 공간 $\mathcal{X} \times \mathcal{Z}$에서 제어된 상태 공간 $\mathcal{X}$로 축소하는 것이다. 저자들은 평균 가치 함수 $\bar{V}(x) := \mathbb{E}_{Z \sim \mu}[V(x, Z)]$를 정의한다.
다음 컨텍스트가 $\mu$로부터 독립적으로 추출되기 때문에, 평균 가치 $\bar{V}^\star$에 대한 벨만 최적 방정식은 $\mu$에 대해 평균화된 전이 커널에만 의존한다:
$$\bar{V}^\star(x) = \mathbb{E}_{Z \sim \mu} \left[ \max_{a \in \mathcal{A}} \left( r(x, a, Z) + \gamma \sum_{x'} P(x' | x, a, Z) \bar{V}^\star(x') \right) \right]$$
이 공식화는 학습자가 $\mathbb{R}^{|\mathcal{X}||\mathcal{Z}|}$가 아닌 $\mathbb{R}^{|\mathcal{X}|}$의 벡터를 추정할 수 있게 하여, 컨텍스트 공간의 크기 $|\mathcal{Z}|$로부터 복잡도를 효과적으로 분리시킨다.

2.2 알고리즘 프레임워크: 분산 감소 할빙(Variance-Reduced Halving)

본 논문은 평균 가치 함수를 추정하기 위해 분산 감소 가치 반복(variance-reduced value iteration) 기법(Sidford et al., 2019에서 변형됨)을 채택한다. 이 알고리즘은 두 단계로 작동한다:

1. 오프라인 단계 (앵커 추정): 알고리즘은 초기 가치 추정치에 적용된 벨만 연산자에 대한 고확률 하한(anchor)을 계산한다. 이는 경험적 평균의 분산을 제어하기 위해 대규모 배치 샘플을 필요로 한다.
2. 온라인/반복 단계 (분산 감소 업데이트): 후속 반복은 현재의 벨만 백업과 앵커 사이의 차이만을 추정한다. 동일한 샘플을 차이 항에 재사용함으로써 분산이 크게 감소하며, 이를 통해 더 작은 샘플 사이즈로도 운용이 가능하다.

알고리즘은 목표 정확도 $\epsilon$에 도달할 때까지 오차 경계 $u$를 반복적으로 절반으로 줄이는 서브루틴(HALFERR)을 사용하는 "할빙(halving)" 전략을 사용한다.

2.3 알려지지 않은 역학 처리

완전히 알려지지 않은 정칙에서는 알고리즘이 샘플링된 컨텍스트를 가지고 있더라도 벨만 백업을 정확히 계산할 수 없다. 대신, 모든 샘플링된 컨텍스트 $Z$에 대해, 알고리즘은 $X' \sim P(\cdot | x, \pi(x, Z), Z)$로부터 전이 샘플을 추출하여 백업의 불편 추정치(unbiased estimate)를 형성한다. 이는 컨텍스트와 전이의 무작위성을 모두 고려하는 마팅게일 분해(martingale decomposition)를 통해 엄격하게 경계 지어지는 추가적인 분산 항을 도입한다.

3. 주요 결과

3.1 알려진 역학, 알려지지 않은 컨텍스트 분포 ($P, r$은 알고, $\mu$는 모름)

전이 커널과 보상이 알려진 경우, 본 논문은 샘플 복잡도가 컨텍스트 공간의 크기 $|\mathcal{Z}|$와 무관함을 입증한다.

• 복잡도: 알고리즘은 $(n_{\text{learn}}, m_{\text{query}}) = (\tilde{O}(\beta^3/\epsilon^2), 0)$의 복잡도를 달성한다. 여기서 $\beta = (1-\gamma)^{-1}$이다.
• 최적성: $\Omega(\beta^3/\epsilon^2)$의 미니맥스 하한이 증명되었으며, 이는 상한이 로그 인자를 제외하고 타이트함을 보여준다.
• 시사점: 학습자는 모든 $(x, z)$ 쌍에 대한 가치를 추정할 필요가 없으며, 평균 가치 $\bar{V}$를 추정하는 것만으로 충분하다. 쿼리 단계에서는 추가 샘플이 필요하지 않다($m_{\text{query}}=0$). 왜냐นั้น $\bar{V}$에서 $V(x, z)$로의 리프트(lift)는 알려진 $P$와 $r$을 사용하는 결정론적 벨만 백업이기 때문이다.

3.2 완벽한 1단계 룩어헤드(Perfect One-Step Lookahead)에 대한 적용

이 프레임워크는 완벽한 1단계 전이 룩어헤드 설정(에이전트가 행동하기 전에 가능한 모든 행동에 대한 모든 가능한 다음 상태를 관찰하는 설정)에 적용된다.

• 이 설정은 컨텍스트 공간이 가능한 모든 전이 텐서의 집합인 외생적 컨텍스트 MDP의 특수한 사례임을 보여준다.
• 결과적으로 $\tilde{O}(\beta^3/\epsilon^2)$의 룩어헤드 샘플 복잡도를 도출한다.
• 개별 테이블형 전이 관점에서 이는 $\tilde{O}(\beta^3 |\mathcal{X}||\mathcal{A}|/\epsilon^2)$에 해당하며, Lu et al. (2025)의 기존 바운드를 $\beta$만큼 개선한다.

3.3 완전히 알려지지 않은 정칙 ($P$와 $\mu$ 모두 모름)

역학(dynamics)과 컨텍스트 분포가 모두 알려지지 않은 경우, 본 논문은 정책 평가(PE)에 집중한다.

• 복잡도: 알고리즘은 $(n_{\text{learn}}, m_{\text{query}}) = (\tilde{O}(|\mathcal{X}|\beta^3/\epsilon^2), \tilde{O}(\beta^2/\epsilon^2))$를 달성한다.
• 트레이드오프: 오프라인 단계에서는 평균 가치 $\bar{V}^\pi$를 학습하며(컨텍스트 샘플당 $|\mathcal{X}|$번의 전이 소요), 쿼리 단계에서는 실현된 $(x, z)$에 대해 신선한 샘플을 사용하여 최종 1단계 전이 기댓값을 추정한다.
• 하한: 본 논문은 다음을 입증하는 일치하는 하한을 증명한다:
  • 의사결정 시점의 정확도를 위해 $m_{\text{query}} = \Omega(\beta^2/\epsilon^2)$가 필요하다.
  • 오프라인 학습을 위해 $n_{\text{learn}} + |\mathcal{X}| m_{\text{query}} = \Omega(|\mathcal{X}|\beta^3/\epsilon^2)$가 필요하다.
• 의의: 완전히 알려지지 않은 정칙에서도 복잡도는 $|\mathcal{Z}|$와 무관하게 유지된다. 알려지지 않은 역학은 확장된 공간에서의 추정을 강제하는 것이 아니라, $\bar{V}$에서 $V(x, z)$로의 의사결정 시점 "리프트"에 확률성을 도입할 뿐이다.

4. 의의 및 주장

본 논문은 외생적 컨텍스트 MDP 학습을 위한 최초의 미니맥스 최적(minimax-optimal) PAC 바운드를 제공한다고 주장한다. 주요 기여는 다음과 같다:

1. 컨텍스트 크기와의 분리: 컨텍스트 공간이 커짐에 따른 통계적 비용을 피할 수 있음을 보여준다. 샘플 복잡도는 컨텍스트 공간 $|\mathcal{Z}|$가 아닌 제어된 상태 공간 $|\mathcal{X}|$에 의존한다.
2. 최적 비율: 도출된 비율($알려진 역학의 경우$\tilde{O}(\beta^3/\epsilon^2)$, 완전히 알려지지 않은 경우$\tilde{O}(|\mathcal{X}|\beta^3/\epsilon^2)$)은 로그 인자를 제외하고 미니맥스 최적임이 입증되었다.
3. 통합된 프레임워크: 이 프레임워크는 완벽한 1단계 룩어헤드의 기존 결과들을 포괄하고 개선하며, 오프라인/온라인 샘플 트레이드오프에 대한 엄격한 분석을 제공한다.

저자들은 순수 온라인 지점($n_{\text{learn}}=0$)과 재사용 가능한 오라클 지점(그들이 제안한 알고리즘)에서는 하한이 타이트하지만, 중간 단계의 트레이드오프에 대한 타이트함은 여전히 미해결 과제로 남아 있다고 언급한다. 또한, 전이 샘플로부터 행동에 대한 최대값을 추정할 때 발생하는 선택 편향(selection bias)을 다루는 어려움 때문에, 완전히 알려지지 않은 정칙에서의 최적 정책 추출(BPE)로의 확장은 향후 연구 과제로 남겨두었다.