Distributional Hamilton-Jacobi-Bellman Equations for Continuous-Time Reinforcement Learning
Continuous-time reinforcement learning offers an appealing formalism for
describing control problems in which the passage of time is not naturally
divided into discrete increments. Here we consider the problem of predicting
the distribution of returns obtained by an agent interacting in a
continuous-time, stochastic environment. Accurate return predictions have
proven useful for determining optimal policies for risk-sensitive control,
learning state representations, multiagent coordination, and more. We begin by
establishing the distributional analogue of the Hamilton-Jacobi-Bellman (HJB)
equation for It\^o diffusions and the broader class of Feller-Dynkin processes.
We then specialize this equation to the setting in which the return
distribution is approximated by $N$ uniformly-weighted particles, a common
design choice in distributional algorithms. Our derivation highlights
additional terms due to statistical diffusivity which arise from the proper
handling of distributions in the continuous-time setting. Based on this, we
propose a tractable algorithm for approximately solving the distributional HJB
based on a JKO scheme, which can be implemented in an online control algorithm.
We demonstrate the effectiveness of such an algorithm in a synthetic control
problem.