HJB 방정식: 연속시간 강화학습과 디퓨전 모델이 만나는 지점

강화학습(RL)이나 디퓨전 모델을 공부하다 보면 “결국 다 비슷한 수학으로 귀결된다”는 말을 종종 들어요.
그 중심에 HJB(Hamilton–Jacobi–Bellman) 방정식이 있고, 이 한 줄이 연속시간 RL과 디퓨전 생성을 같은 언어로 묶어줍니다.

1) Bellman 방정식에서 HJB로: “연속시간”이 왜 중요한가

요약하면, 이산시간 Bellman 최적방정식을 시간 간격 (h)로 두고 (h \downarrow 0) 극한을 보내면 PDE(편미분방정식) 형태의 HJB가 됩니다.
이 과정은 “지금 당장 행동의 이득 + 미래 가치”라는 Bellman의 로컬 최적화 논리를 그대로 유지하되, 시간이 촘촘해져서 방정식 형태가 달라지는 거예요.

특히 노이즈가 없는 결정론적 시스템에서는, 가치함수 (V(t,x))가 (-\partial_t V = H(t,x,\nabla V)) 꼴을 만족하고, 여기서 Hamiltonian (H)는 “어떤 행동을 하는 게 제일 이득인가”를 한 번에 모아놓은 함수입니다.
이게 중요한 이유는 실제 세계(로봇, 금융, 물리 시뮬레이션)는 이산 스텝보다 연속시간 모델로 표현되는 경우가 많기 때문이에요.

2) HJB = Hamilton–Jacobi: RL이 물리학 수학과 같아지는 순간

요약하면, Bellman이 연속시간으로 확장하면서 얻은 PDE가 19세기 고전역학의 Hamilton–Jacobi 방정식과 구조적으로 동일하다는 점이 핵심입니다.
즉 (S \leftrightarrow V)(작용/Action 함수와 가치함수의 대응), (q \leftrightarrow x)(좌표와 상태의 대응)만 놓으면 “같은 방정식”이 됩니다.

원문에서는 보상을 라그랑지안(Lagrangian)으로 바꿔 쓰는 트릭, 즉 **(r=-L)**로 두고 Hamiltonian을 **(\sup_a{p^\top f – L})**로 정의하면 Hamilton–Jacobi와 거의 같은 꼴이 된다고 설명해요.
이 연결이 갖는 의미는 큽니다. RL을 단지 “보상 최대화”로만 보지 않고, **최적제어(optimal control)**나 물리 기반 관점으로 해석할 길이 열리거든요.

현업 시나리오로는, 연속시간 기반의 제어 문제(예: 드론 자세 제어, 공정 제어)를 RL로 풀 때 ‘내가 푸는 문제가 본질적으로 PDE/최적제어다’라는 감각이 설계 선택(모델 기반 vs 모델 프리, 안정성 제약 등)을 더 탄탄하게 해줘요.

3) 노이즈가 들어오면 “제너레이터”가 등장한다: Controlled Diffusion과 HJB

요약하면, 연속시간 + 확률 노이즈가 들어오면 상태는 SDE(확률미분방정식, Itô process)로 움직이고, HJB에는 **(\mathcal L^a)**라는 **infinitesimal generator(무한소 생성자)**가 들어갑니다.
직관적으로는 “아주 짧은 시간 뒤 가치함수 기대값이 어떻게 변하는지”를 1차로 요약한 연산자예요.

상태가
[
dX_t = f(X_t,a_t)dt + \Sigma(X_t,a_t)dW_t
]
로 움직일 때, generator는

• 드리프트 항: (\nabla V^\top f)
• 확산(곡률) 항: (\tfrac12 \mathrm{Tr}(\Sigma\Sigma^\top \nabla^2 V))
  로 구성됩니다.

그래서 할인율 (\rho)가 있는 무한시간 문제의 HJB는
[
\rho V(x)=\max_a{r(x,a)+\mathcal L^aV(x)}
]
가 되고, 여기서 “노이즈가 없으면 사라지는” (\nabla^2V) 항이 생겨요. 이 부분이 연속시간 확률 제어의 난이도이자, 동시에 디퓨전 모델과 연결되는 출발점이기도 합니다.

4) 연속시간 RL의 핵심 루프: Policy Iteration을 신경망으로 돌리는 방법

요약하면, 원문은 HJB를 푸는 실용적인 방법으로 **Policy Iteration(정책 반복)**을 제시하고, **가치함수 (V_\theta)**와 **정책 (\alpha_\phi)**를 MLP로 근사해요.
구조는 우리가 아는 RL과 같고, 다만 “Bellman backup”이 테이블이 아니라 **HJB 잔차(residual)**와 generator 계산으로 바뀝니다.

핵심 단계는 아래 3개예요.

• Policy evaluation(정책 평가): 현재 정책 (\alpha_k)에 대한 (V_k)를 맞춘다
  • 원문은 Feynman–Kac 표현을 이용해, 시뮬레이션 롤아웃(몬테카를로)로 적분 형태의 가치를 추정합니다.
  • 이 방식은 “PDE를 직접 풀기 어려우니, 확률 과정 샘플링으로 값을 맞춘다”는 실용적 선택이에요.
• Policy improvement(정책 개선):
  [
  \alpha_{k+1}(x)\in\arg\max_a{r(x,a)+\mathcal L^aV_k(x)}
  ]
  • 즉 (Q)를 정의해서 greedy 업데이트를 합니다. 연속시간에서의 Q는
    [
    Q(x,a)=\frac{1}{\rho}(r+\mathcal L^aV)
    ]
    로 두고 **(V=\max_a Q)**가 성립해요.
• Diagnosis(진단): HJB residual
  [
  \mathcal R_{\mathrm{HJB}}=\rho V-\max_a{r+\mathcal L^aV}
  ]
  이 더 이상 줄지 않으면 멈춥니다.

여기서 구현 포인트는 generator에 필요한 (\nabla V), (\nabla^2V)를 autograd로 구한다는 점이에요. 즉 연속시간 RL은 “미분 가능한 가치함수 근사”와 궁합이 좋다는 결론이 자연스럽게 나옵니다.

5) 모델 프리로 가면: 연속시간 Q-learning과 TD 타깃의 형태

요약하면, 정책 반복은 **모델 기반(model-based)**이라 (f,\Sigma) 즉 dynamics를 알아야 해요.
반면 원문은 연속시간 Q-learning을 소개하면서, 짧은 전이 ((X_t,a_t,r_t,X_{t+\Delta t}))만으로 학습하는 길도 보여줍니다.

핵심은 (\Delta t)가 작은 구간에서 TD 타깃을

• (r_t\Delta t) (연속시간이므로 보상이 “시간 적분”이라 스케일이 붙음)
• (e^{-\rho \Delta t}\bar V(X_{t+\Delta t})) (할인 연속버전)
  로 잡는다는 점이에요.

이렇게 하면 우리가 익숙한 actor-critic처럼

• critic: (Q_\psi)를 MSE로 맞추고
• actor: (Q_\psi(X_t,a_\omega(X_t)))를 키우도록 업데이트
  하는 형태로 이어집니다.
  현업 감각으로는, 시뮬레이터가 없거나 dynamics를 모르지만, 아주 촘촘한 로그/샘플이 있는 시스템(예: HFT, 센서 스트림 제어)에서 이런 형태가 특히 매력적이에요.

6) 디퓨전 모델을 최적제어로 보면: “역과정 샘플링 = 최적 정책”이 된다

요약하면, 원문에서 가장 흥미로운 연결은 디퓨전 모델의 reverse-time 과정을 유한시간 확률 최적제어 문제로 재해석하는 부분이에요.
여기서 가치함수를
[
V(x,t) = -\log p_{T-t}(x)
]
(역시간 마진널의 negative log-density)로 두면, 이 (V)가 유한시간 HJB를 만족하도록 정리됩니다.

이때 제어 (u(x,t))를 드리프트에 더해
[
dX_t^u = (\mu + \Sigma u)dt + \Sigma dW_t
]
형태로 만들고, 어떤 (u)를 선택해야 최종 분포가 (p_{\text{data}})가 되는지 보는 거예요.
그리고 최적 제어는
[
u^*(x,t)= -\Sigma^\top\nabla V = \Sigma^\top \nabla \log p_{T-t}(x)
]
가 되어, 결국 우리가 디퓨전에서 말하는 **score(점수 함수, (\nabla \log p))**가 그대로 튀어나옵니다.

정리하면, 디퓨전 모델에서 네트워크가 하는 일은 관점에 따라

• **score (s(x,t))**를 맞추는 일일 수도 있고
• **스케일된 제어 (u(x,t))**를 맞추는 일일 수도 있어요.
  그리고 “정답 제어와의 오차”가 **제곱 형태(Quadratic)**로 비용에 들어가면서, 실무에서 쓰는 denoising score matching 목적함수와 이어진다는 해석이 가능합니다.

마무리: HJB를 “공식”이 아니라 “연결고리”로 써보면 좋겠어요

이 글의 요지는 단순히 HJB를 소개하는 게 아니라, Bellman → 연속시간 → 확률 제어 → 디퓨전이 한 줄로 이어진다는 점이에요.
강화학습을 하다가 디퓨전을 만나도, 디퓨전을 하다가 최적제어를 만나도, 결국 같은 구조를 다시 보게 됩니다.

다음에 할 일로는, 본인이 다루는 문제를 한 번 점검해보면 좋아요.
“내 환경은 이산시간인가, 사실상 연속시간인가?”, “노이즈가 본질적인가?”, “모델 기반( generator 계산 가능) 접근이 가능한가?” 같은 질문이 알고리즘 선택을 훨씬 명확하게 만들어줄 거예요.