Yet Another Distributional Bellman Equation

TL;DR

This paper introduces a generalized distributional Bellman equation for Markov Decision Processes that considers complex functionals of joint distributions, extending traditional expectation-based approaches to more general distributional objectives.

Contribution

It develops a novel framework defining a lifted MDP with a distribution-based state space and derives a distributional Bellman equation, unifying various existing models.

Findings

Derives a distributional Bellman equation applicable to general functionals.

Shows standard MDPs and quantile MDPs as special cases.

Applies the framework to an optimal transport problem variant.

Abstract

We consider non-standard Markov Decision Processes (MDPs) where the target function is not only a simple expectation of the accumulated reward. Instead, we consider rather general functionals of the joint distribution of terminal state and accumulated reward which have to be optimized. For finite state and compact action space, we show how to solve these problems by defining a lifted MDP whose state space is the space of distributions over the true states of the process. We derive a Bellman equation in this setting, which can be considered as a distributional Bellman equation. Well-known cases like the standard MDP and quantile MDPs are shown to be special examples of our framework. We also apply our model to a variant of an optimal transport problem.