Adaptive Partitioning and Learning for Stochastic Control of Diffusion Processes

We study reinforcement learning for controlled diffusion processes with unbounded continuous state spaces, bounded continuous actions, and polynomially growing rewards: settings that arise naturally in finance, economics, and operations research. To overcome the challenges of continuous and high-dimensional domains, we introduce a model-based algorithm that adaptively partitions the joint state-action space. The algorithm maintains estimators of drift, volatility, and rewards within each partition, refining the discretization whenever estimation bias exceeds statistical confidence. This adaptive scheme balances exploration and approximation, enabling efficient learning in unbounded domains. Our analysis establishes regret bounds that depend on the problem horizon, state dimension, reward growth order, and a newly defined notion of zooming dimension tailored to unbounded diffusion processes. The bounds recover existing results for bounded settings as a special case, while extending theoretical guarantees to a broader class of diffusion-type problems. Finally, we validate the effectiveness of our approach through numerical experiments, including applications to high-dimensional problems such as multi-asset mean-variance portfolio selection.

💡 Research Summary

This paper addresses reinforcement learning (RL) for stochastic control problems governed by diffusion processes that feature unbounded continuous state spaces, bounded continuous action spaces, and rewards that may grow polynomially with the state. Such settings naturally arise in finance (e.g., portfolio optimization), economics, and operations research, yet they have received limited theoretical treatment in the RL literature, which traditionally assumes either finite or bounded continuous domains and bounded rewards.

Problem formulation. The authors consider a discrete‑time Markov decision process (MDP) that approximates a continuous‑time diffusion. At each step (h) the state evolves as
(X_{h+1}=X_h+\mu_h(X_h,A_h)\Delta+\sigma_h(X_h,A_h)B_h\sqrt{\Delta}),
where (B_h\sim\mathcal N(0,I_{d_S})). The action space (A\subset\mathbb R^{d_A}) is a closed hyper‑cube, and the instantaneous reward distribution (R_h) has mean (\bar R_h(x,a)) that may scale like (|x|^{m}) for some integer (m\ge 0). The goal is to learn a policy (\pi) that maximizes the expected cumulative reward over a finite horizon (H).

Algorithmic contribution. The core of the work is a model‑based, adaptive‑partitioning algorithm. The joint state‑action space is initially covered by a coarse grid of hyper‑rectangular cells. Within each cell the algorithm maintains empirical estimators of the drift (\mu_h), volatility (\sigma_h), and mean reward (\bar R_h). A confidence interval is constructed for each estimator based on the number of samples collected in that cell. When the estimated bias of any quantity exceeds its confidence radius, the corresponding cell is split uniformly into two (or more) sub‑cells. This refinement rule ensures that regions visited frequently (or where the dynamics are highly non‑linear) are represented with finer resolution, while rarely visited regions remain coarse, thereby controlling sample complexity.

The estimated Q‑function is built from the learned drift and volatility via a dynamic‑programming recursion, and actions are selected using an upper‑confidence‑bound (UCB) principle: the algorithm chooses the action that maximizes the sum of the point estimate and a bonus term proportional to the cell’s confidence radius. This yields a natural exploration‑exploitation balance without requiring separate exploration schedules.

Zooming dimension for unbounded domains. A major theoretical novelty is the definition of a “zooming dimension” (z_{\max,c}) that captures the intrinsic difficulty of the problem when the state space is unbounded. Classical zooming dimension (Kleinberg et al., 2008) assumes a bounded metric space and counts how many balls of a given radius are needed to cover the region where the optimal value function varies significantly. In the unbounded setting, the authors introduce a weighted volume measure that incorporates the polynomial growth of the reward: each cell (B) receives weight (w(B)=\int_B (1+|x|^{m})dx). The zooming dimension is then the smallest exponent (\alpha) such that (\sum_{B} w(B),\operatorname{diam}(B)^{\alpha}) is bounded. Intuitively, (z_{\max,c}) reflects how “benign” the problem is; for many financial models it is far smaller than the ambient dimension (d_S+d_A).

Theoretical analysis. The regret analysis proceeds through four technical steps:

1. Concentration for drift and volatility. Using only Lipschitz continuity of (\mu_h) and (\sigma_h), the authors derive matrix‑valued concentration inequalities for the empirical covariance matrices. They combine a matrix Azuma inequality with a Bernstein‑type bound to handle the fact that volatility estimators involve products of Gaussian increments.

2. Reward estimation under polynomial growth. Because (\bar R_h(x,a)) can be unbounded, the bias and variance of the reward estimator scale with (\operatorname{diam}(B)^{m}) and (\operatorname{diam}(B)^{2m}), respectively. The analysis shows that these terms remain controlled as long as the cell diameters shrink appropriately during refinement.

3. Bounding the number of cells. The refinement rule guarantees that the total number of cells after (K) episodes grows at most like (\tilde O\big(K^{(z_{\max,c}+1)/(z_{\max,c}+2)}\big)). This sub‑linear growth is crucial for keeping the computational burden manageable.

4. Regret bound. Combining the above ingredients, the authors prove that the cumulative regret after (K) episodes satisfies
   \